Spaces:
Build error
Build error
File size: 2,685 Bytes
b1fc48a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 |
import streamlit as st
import torch
import torchvision.transforms as transforms
# from torch import nn
# import cv2
# import numpy as np
# import torch
# import torchvision
# from torchvision import models, transforms
# from torch.utils.data import DataLoader
# from torch.utils.data.dataset import Dataset
# import os
# import matplotlib.pyplot as plt
# import face_recognition
# import sys
# import time
# from torch.autograd import Variable
from function import Model, validation_dataset, predict
# Set Streamlit page config
st.set_page_config(
page_title="DeepFake Detection",
layout="centered",
page_icon=":mask:"
)
# Load the pre-trained model
@st.cache_resource
def load_model():
model = Model(num_classes=2).cuda()
model_path = "model.pt" # Update with actual model path
model.load_state_dict(torch.load(model_path))
model.eval()
return model
model = load_model()
st.title("DeepFake Detection App")
st.divider()
st.sidebar.header("Upload Video")
uploaded_video = st.sidebar.file_uploader("Choose a video file", type=["mp4", "avi", "mov"])
if uploaded_video:
st.sidebar.write("Video uploaded successfully!")
# Save the uploaded video locally for processing
with open("uploaded_video.mp4", "wb") as f:
f.write(uploaded_video.getbuffer())
st.video("uploaded_video.mp4")
# Preprocess and analyze the video
with st.spinner("Processing video..."):
try:
# Validation dataset
video_dataset = validation_dataset(
video_names=["uploaded_video.mp4"],
sequence_length=20,
transform=transforms.Compose([
transforms.ToPILImage(),
transforms.Resize((112, 112)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
)
video_frames = video_dataset[0] # Extract frames from video
if video_frames is not None:
prediction = predict(model, video_frames)
st.subheader("Prediction: ")
st.write(f"{'REAL' if prediction[0] == 1 else 'FAKE'}")
st.subheader("Confidence: ")
st.write(f"{prediction[1]:.2f}%")
else:
st.error("No faces detected in the video.")
except Exception as e:
st.error(f"An error occurred during processing: {e}")
else:
st.sidebar.write("Upload a video file.")
|