Spaces:
Runtime error
Runtime error
import torch.nn.functional as F | |
import torch | |
import torch.nn as nn | |
import torch.optim as optim | |
from torch.utils.data import DataLoader | |
from sklearn.metrics import accuracy_score, precision_recall_fscore_support | |
from torch.optim.lr_scheduler import CosineAnnealingLR | |
from tqdm import tqdm | |
import warnings | |
warnings.filterwarnings("ignore") | |
import cv2 | |
import numpy as np | |
import matplotlib.pyplot as plt | |
import pywt | |
from utils.config import cfg | |
from dataset.real_n_fake_dataloader import Extracted_Frames_Dataset | |
from utils.data_transforms import get_transforms_train, get_transforms_val | |
from net.Multimodalmodel import Image_n_DCT | |
import os | |
import json | |
import torch | |
from torchvision import transforms | |
from torch.utils.data import DataLoader, Dataset | |
from PIL import Image | |
import numpy as np | |
import pandas as pd | |
import cv2 | |
import argparse | |
class Test_Dataset(Dataset): | |
def __init__(self, test_data_path = None, transform = None, image_path = None, multi_modal = "dct"): | |
""" | |
Args: | |
returns: | |
""" | |
self.multi_modal = multi_modal | |
if test_data_path is None and image_path is not None: | |
self.dataset = [[image_path, 2]] | |
self.transform = transform | |
else: | |
self.transform = transform | |
self.real_data = os.listdir(test_data_path + "/real") | |
self.fake_data = os.listdir(test_data_path + "/fake") | |
self.dataset = [] | |
for image in self.real_data: | |
self.dataset.append([test_data_path + "/real/" + image, 1]) | |
for image in self.fake_data: | |
self.dataset.append([test_data_path + "/fake/" + image, 0]) | |
def __len__(self): | |
return len(self.dataset) | |
def __getitem__(self, idx): | |
sample_input = self.get_sample_input(idx) | |
return sample_input | |
def get_sample_input(self, idx): | |
rgb_image = self.get_rgb_image(idx) | |
label = self.get_label(idx) | |
if self.multi_modal == "dct": | |
dct_image = self.get_dct_image(idx) | |
sample_input = {"rgb_image": rgb_image, "dct_image": dct_image, "label": label} | |
# dct_image = self.get_dct_image(idx) | |
elif self.multi_modal == "fft": | |
fft_image = self.get_fft_image(idx) | |
sample_input = {"rgb_image": rgb_image, "dct_image": fft_image, "label": label} | |
elif self.multi_modal == "hh": | |
hh_image = self.get_hh_image(idx) | |
sample_input = {"rgb_image": rgb_image, "dct_image": hh_image, "label": label} | |
else: | |
AssertionError("multi_modal must be one of (dct:discrete cosine transform, fft: fast forier transform, hh)") | |
return sample_input | |
def get_fft_image(self, idx): | |
gray_image_path = self.dataset[idx][0] | |
gray_image = cv2.imread(gray_image_path, cv2.IMREAD_GRAYSCALE) | |
fft_image = self.compute_fft(gray_image) | |
if self.transform: | |
fft_image = self.transform(fft_image) | |
return fft_image | |
def compute_fft(self, image): | |
f = np.fft.fft2(image) | |
fshift = np.fft.fftshift(f) | |
magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1) # Add 1 to avoid log(0) | |
return magnitude_spectrum | |
def get_hh_image(self, idx): | |
gray_image_path = self.dataset[idx][0] | |
gray_image = cv2.imread(gray_image_path, cv2.IMREAD_GRAYSCALE) | |
hh_image = self.compute_hh(gray_image) | |
if self.transform: | |
hh_image = self.transform(hh_image) | |
return hh_image | |
def compute_hh(self, image): | |
coeffs2 = pywt.dwt2(image, 'haar') | |
LL, (LH, HL, HH) = coeffs2 | |
return HH | |
def get_rgb_image(self, idx): | |
rgb_image_path = self.dataset[idx][0] | |
rgb_image = Image.open(rgb_image_path) | |
if self.transform: | |
rgb_image = self.transform(rgb_image) | |
return rgb_image | |
def get_dct_image(self, idx): | |
rgb_image_path = self.dataset[idx][0] | |
rgb_image = cv2.imread(rgb_image_path) | |
dct_image = self.compute_dct_color(rgb_image) | |
if self.transform: | |
dct_image = self.transform(dct_image) | |
return dct_image | |
def get_label(self, idx): | |
return self.dataset[idx][1] | |
def compute_dct_color(self, image): | |
image_float = np.float32(image) | |
dct_image = np.zeros_like(image_float) | |
for i in range(3): | |
dct_image[:, :, i] = cv2.dct(image_float[:, :, i]) | |
return dct_image | |
class Test: | |
def __init__(self, model_path, multi_modal = "dct"): | |
self.model_path = model_path | |
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
print(self.device) | |
# Load the model | |
self.model = Image_n_DCT() | |
self.model.load_state_dict(torch.load(self.model_path, map_location = self.device)) | |
self.model.to(self.device) | |
self.model.eval() | |
self.multi_modal = multi_modal | |
def testimage(self, image_path): | |
test_dataset = Test_Dataset(transform = get_transforms_val(), image_path = image_path, multi_modal = self.multi_modal) | |
inputs = test_dataset[0] | |
rgb_image, dct_image = inputs['rgb_image'].to(self.device), inputs['dct_image'].to(self.device) | |
output = self.model(rgb_image.unsqueeze(0), dct_image.unsqueeze(0)) | |
# print(output.shape) | |
_, predicted = torch.max(output.data, 1) | |
return 'real' if predicted==1 else 'fake' |