Spaces:

22GC22
/

deepsaif

Running

Upload 17 files

2fbad05 verified almost 2 years ago

8.25 kB

import streamlit as st
from PIL import Image
from models import mesonet, mesoinception, fal_detector, local_detector
from utils.visualization import display_results
from utils.preprocessing import preprocess_image, preprocess_video
from utils.preprocessing import preprocess_image, generate_local_image
import numpy as np
import cv2
import tempfile
import os

# Initialize session state for tabs and uploaded files
if "active_tab" not in st.session_state:
    st.session_state["active_tab"] = "Face Photoshop Detection"

if "uploaded_file" not in st.session_state:
    st.session_state["uploaded_file"] = None

# Load models
models = {
    "MesoNet": mesonet.load_mesonet("models/weights/Meso4_DF.h5"),
    "MesoInception": mesoinception.load_mesonetInception("models/weights/MesoInception_DF.h5"),
    "Photoshop FALdetector Global": fal_detector.load_fal_detector("models/weights/global.pth"),
    "Photoshop FALdetector Local": local_detector.load_local_detector("models/weights/local.pth", gpu_id=-1),
}


st.title("DeepSAIF")
# Create tabs for different functionalities
tab1, tab2, tab3 = st.tabs(["Face Photoshop Detection", "DeepFake Detection for Images", "DeepFake Detection for Videos"])

# Tab 1: Photoshop Detection
with tab1:
    if st.session_state["active_tab"] != "Face Photoshop Detection":
        st.session_state["uploaded_file"] = None
        st.session_state["active_tab"] = "Face Photoshop Detection"

    st.header("Face Photoshop Detection")
    uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png"], key="photoshop")

    if uploaded_file:
        st.session_state["uploaded_file"] = uploaded_file
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption="Uploaded Image", use_column_width=True)
        local_image = generate_local_image(image)

        # Run inference on all models
        results = {}
        for model_name, model in models.items():
            if model_name == "Photoshop FALdetector Global":
                results[model_name] = fal_detector.predict_fal_detector(model, image)
            elif model_name == "Photoshop FALdetector Local":
                heatmap_path, prediction = local_detector.predict_and_generate_heatmap(model, image)
                if heatmap_path:
                    # Display the heatmap using Streamlit
                    st.image(heatmap_path, caption=f"Heatmap for {model_name}", use_container_width=True)
                    
                    # Delete the temporary heatmap file after display
                    os.remove(heatmap_path)
                    os.remove('cropped_input.jpg')
                    os.remove('warped.jpg')
                else:
                    st.error(f"Failed to generate heatmap for {model_name}")
                results[model_name] = prediction
            # elif model_name == "Global Classifier":
            #     results[model_name] = global_classifier.classify_fake(model, image)


        # Display results
        display_results(results)

# Tab 2: DeepFake Detection for Images
with tab2:
    if st.session_state["active_tab"] != "DeepFake Detection for Images":
        st.session_state["uploaded_file"] = None
        st.session_state["active_tab"] = "DeepFake Detection for Images"

    st.header("DeepFake Detection for Images")

    uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png"], key="deepfake_image")

    if uploaded_file:
        st.session_state["uploaded_file"] = uploaded_file
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption="Uploaded Image", use_column_width=True)
        local_image = generate_local_image(image)
        # Preprocess the image
        # preprocessed_image = preprocess_image(uploaded_file)

        # Run inference on all models
        results = {}
        for model_name, model in models.items():
            if model_name == "MesoNet":
                results[model_name] = mesonet.predict_mesonet(model, image)
            elif model_name == "MesoInception":
                results[model_name] = mesoinception.predict_mesonetInception(model, image)

        # Display results
        display_results(results)



def confident_strategy(pred, t=0.8):
    """
    Implements the confident averaging strategy for predictions.
    Args:
        pred (list[float]): List of predictions for each frame.
        t (float): Threshold for high-confidence fake detection.
    Returns:
        float: Final confidence score for the video.
    """
    if len(pred) == 0:
        return np.nan
    pred = np.array(pred)
    sz = len(pred)
    fakes = np.count_nonzero(pred > t)
    if fakes > sz // 2.5 and fakes > 11:
        return np.mean(pred[pred > t])
    elif np.count_nonzero(pred < 0.2) > 0.9 * sz:
        return np.mean(pred[pred < 0.2])
    else:
        return np.mean(pred)

# Tab 3: DeepFake Detection for Videos
with tab3:
    if st.session_state["active_tab"] != "DeepFake Detection for Videos":
        st.session_state["uploaded_file"] = None
        st.session_state["active_tab"] = "DeepFake Detection for Videos"

    st.header("DeepFake Detection for Videos")
    uploaded_file = st.file_uploader("Upload Video", type=["mp4", "avi", "mov"], key="deepfake_video")
    if uploaded_file:
        st.session_state["uploaded_file"] = uploaded_file
        with st.spinner("Processing video..."):
            # Save uploaded file to a temporary location
            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file:
                temp_file.write(uploaded_file.getbuffer())
                video_path = temp_file.name

            try:
                # Test video accessibility
                cap = cv2.VideoCapture(video_path)
                if not cap.isOpened():
                    st.error("Failed to open video file.")
                else:
                    st.success("Video file opened successfully!")

                # Extract frames from video
                frames = preprocess_video(video_path, frame_count=32)
                if len(frames) == 0:
                    st.error("Failed to extract frames from the video.")
                else:
                    # st.success(f"Extracted {len(frames)} frames.")
                    # for frame in frames[:5]:  # Display first 5 frames
                    #     st.image(frame, caption="Extracted Frame")

                    # Dictionary to store model predictions
                    model_results = {
                        "MesoNet": [],
                        "Photoshop FALdetector Global": []
                    }

                    # Iterate over frames and make predictions for each model
                    for frame in frames:
                        preprocessed_frame = preprocess_image(frame)  # Preprocess frame
                        local_image = generate_local_image(preprocessed_frame)

                        # Predictions for MesoNet and Photoshop FALdetector Global
                        model_results["MesoNet"].append(
                            mesonet.predict_mesonet(models["MesoNet"], preprocessed_frame)
                        )
                        model_results["Photoshop FALdetector Global"].append(
                            fal_detector.predict_fal_detector(models["Photoshop FALdetector Global"], local_image)
                        )


                    # Apply the confident averaging strategy for each model
                    final_results = {}
                    for model_name, predictions in model_results.items():
                        final_results[model_name] = confident_strategy(predictions)

                    # Display results
                    st.write("### Video Analysis Results")
                    display_results(final_results)

                    # Optionally show detailed frame predictions per model
                    if st.checkbox("Show Detailed Frame Predictions"):
                        for model_name, predictions in model_results.items():
                            st.write(f"### Predictions for {model_name}")
                            st.bar_chart(predictions)


            finally:
                # Clean up temporary file
                os.remove(video_path)

	import streamlit as st
	from PIL import Image
	from models import mesonet, mesoinception, fal_detector, local_detector
	from utils.visualization import display_results
	from utils.preprocessing import preprocess_image, preprocess_video
	from utils.preprocessing import preprocess_image, generate_local_image
	import numpy as np
	import cv2
	import tempfile
	import os

	# Initialize session state for tabs and uploaded files
	if "active_tab" not in st.session_state:
	st.session_state["active_tab"] = "Face Photoshop Detection"

	if "uploaded_file" not in st.session_state:
	st.session_state["uploaded_file"] = None

	# Load models
	models = {
	"MesoNet": mesonet.load_mesonet("models/weights/Meso4_DF.h5"),
	"MesoInception": mesoinception.load_mesonetInception("models/weights/MesoInception_DF.h5"),
	"Photoshop FALdetector Global": fal_detector.load_fal_detector("models/weights/global.pth"),
	"Photoshop FALdetector Local": local_detector.load_local_detector("models/weights/local.pth", gpu_id=-1),
	}


	st.title("DeepSAIF")
	# Create tabs for different functionalities
	tab1, tab2, tab3 = st.tabs(["Face Photoshop Detection", "DeepFake Detection for Images", "DeepFake Detection for Videos"])

	# Tab 1: Photoshop Detection
	with tab1:
	if st.session_state["active_tab"] != "Face Photoshop Detection":
	st.session_state["uploaded_file"] = None
	st.session_state["active_tab"] = "Face Photoshop Detection"

	st.header("Face Photoshop Detection")
	uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png"], key="photoshop")

	if uploaded_file:
	st.session_state["uploaded_file"] = uploaded_file
	image = Image.open(uploaded_file).convert("RGB")
	st.image(image, caption="Uploaded Image", use_column_width=True)
	local_image = generate_local_image(image)

	# Run inference on all models
	results = {}
	for model_name, model in models.items():
	if model_name == "Photoshop FALdetector Global":
	results[model_name] = fal_detector.predict_fal_detector(model, image)
	elif model_name == "Photoshop FALdetector Local":
	heatmap_path, prediction = local_detector.predict_and_generate_heatmap(model, image)
	if heatmap_path:
	# Display the heatmap using Streamlit
	st.image(heatmap_path, caption=f"Heatmap for {model_name}", use_container_width=True)

	# Delete the temporary heatmap file after display
	os.remove(heatmap_path)
	os.remove('cropped_input.jpg')
	os.remove('warped.jpg')
	else:
	st.error(f"Failed to generate heatmap for {model_name}")
	results[model_name] = prediction
	# elif model_name == "Global Classifier":
	# results[model_name] = global_classifier.classify_fake(model, image)


	# Display results
	display_results(results)

	# Tab 2: DeepFake Detection for Images
	with tab2:
	if st.session_state["active_tab"] != "DeepFake Detection for Images":
	st.session_state["uploaded_file"] = None
	st.session_state["active_tab"] = "DeepFake Detection for Images"

	st.header("DeepFake Detection for Images")

	uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "png"], key="deepfake_image")

	if uploaded_file:
	st.session_state["uploaded_file"] = uploaded_file
	image = Image.open(uploaded_file).convert("RGB")
	st.image(image, caption="Uploaded Image", use_column_width=True)
	local_image = generate_local_image(image)
	# Preprocess the image
	# preprocessed_image = preprocess_image(uploaded_file)

	# Run inference on all models
	results = {}
	for model_name, model in models.items():
	if model_name == "MesoNet":
	results[model_name] = mesonet.predict_mesonet(model, image)
	elif model_name == "MesoInception":
	results[model_name] = mesoinception.predict_mesonetInception(model, image)

	# Display results
	display_results(results)



	def confident_strategy(pred, t=0.8):
	"""
	Implements the confident averaging strategy for predictions.
	Args:
	pred (list[float]): List of predictions for each frame.
	t (float): Threshold for high-confidence fake detection.
	Returns:
	float: Final confidence score for the video.
	"""
	if len(pred) == 0:
	return np.nan
	pred = np.array(pred)
	sz = len(pred)
	fakes = np.count_nonzero(pred > t)
	if fakes > sz // 2.5 and fakes > 11:
	return np.mean(pred[pred > t])
	elif np.count_nonzero(pred < 0.2) > 0.9 * sz:
	return np.mean(pred[pred < 0.2])
	else:
	return np.mean(pred)

	# Tab 3: DeepFake Detection for Videos
	with tab3:
	if st.session_state["active_tab"] != "DeepFake Detection for Videos":
	st.session_state["uploaded_file"] = None
	st.session_state["active_tab"] = "DeepFake Detection for Videos"

	st.header("DeepFake Detection for Videos")
	uploaded_file = st.file_uploader("Upload Video", type=["mp4", "avi", "mov"], key="deepfake_video")
	if uploaded_file:
	st.session_state["uploaded_file"] = uploaded_file
	with st.spinner("Processing video..."):
	# Save uploaded file to a temporary location
	with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file:
	temp_file.write(uploaded_file.getbuffer())
	video_path = temp_file.name

	try:
	# Test video accessibility
	cap = cv2.VideoCapture(video_path)
	if not cap.isOpened():
	st.error("Failed to open video file.")
	else:
	st.success("Video file opened successfully!")

	# Extract frames from video
	frames = preprocess_video(video_path, frame_count=32)
	if len(frames) == 0:
	st.error("Failed to extract frames from the video.")
	else:
	# st.success(f"Extracted {len(frames)} frames.")
	# for frame in frames[:5]: # Display first 5 frames
	# st.image(frame, caption="Extracted Frame")

	# Dictionary to store model predictions
	model_results = {
	"MesoNet": [],
	"Photoshop FALdetector Global": []
	}

	# Iterate over frames and make predictions for each model
	for frame in frames:
	preprocessed_frame = preprocess_image(frame) # Preprocess frame
	local_image = generate_local_image(preprocessed_frame)

	# Predictions for MesoNet and Photoshop FALdetector Global
	model_results["MesoNet"].append(
	mesonet.predict_mesonet(models["MesoNet"], preprocessed_frame)
	)
	model_results["Photoshop FALdetector Global"].append(
	fal_detector.predict_fal_detector(models["Photoshop FALdetector Global"], local_image)
	)


	# Apply the confident averaging strategy for each model
	final_results = {}
	for model_name, predictions in model_results.items():
	final_results[model_name] = confident_strategy(predictions)

	# Display results
	st.write("### Video Analysis Results")
	display_results(final_results)

	# Optionally show detailed frame predictions per model
	if st.checkbox("Show Detailed Frame Predictions"):
	for model_name, predictions in model_results.items():
	st.write(f"### Predictions for {model_name}")
	st.bar_chart(predictions)


	finally:
	# Clean up temporary file
	os.remove(video_path)