File size: 1,539 Bytes
627e429 |
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 |
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" # to change if you have a GPU with at least 12Go RAM (it will save you a lot of time !)
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" # to change if you have a GPU with at least 12Go RAM (it will save you a lot of time !)
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))
|