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))