|
import os |
|
from matplotlib import pyplot as plt |
|
from MantraNet.mantranet import pre_trained_model, check_forgery |
|
from BusterNet.BusterNetCore import create_BusterNet_testing_model |
|
from BusterNet.BusterNetUtils import simple_cmfd_decoder, visualize_result |
|
import streamlit as st |
|
import cv2 |
|
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" |
|
|
|
st.header("IMD Demo") |
|
device = "cpu" |
|
|
|
def check_image_buster(img_path): |
|
busterNetModel = create_BusterNet_testing_model( 'BusterNet/pretrained_busterNet.hd5' ) |
|
rgb = cv2.imread(img_path) |
|
pred = simple_cmfd_decoder( busterNetModel, rgb ) |
|
figure = visualize_result( rgb, pred, pred, figsize=(20,20), title='BusterNet CMFD') |
|
st.pyplot(figure) |
|
|
|
def check_image_mantra(img_path): |
|
device = "cpu" |
|
MantraNetmodel = pre_trained_model( |
|
weight_path="MantraNet/MantraNetv4.pt", device=device |
|
) |
|
fig = check_forgery(MantraNetmodel, img_path=img_path, device=device) |
|
st.pyplot(fig) |
|
|
|
|
|
uploaded_image = st.file_uploader("Upload your image", type=["jpg", "png","jpeg"]) |
|
if uploaded_image is not None: |
|
with open(os.path.join("images", uploaded_image.name), "wb") as f: |
|
f.write(uploaded_image.read()) |
|
st.write("BusterNet") |
|
check_image_buster(os.path.join("images", uploaded_image.name)) |
|
st.write("MantraNet") |
|
check_image_mantra(os.path.join("images", uploaded_image.name)) |
|
|
