#######################################################################################
#
# MIT License
#
# Copyright (c) [2025] [leonelhs@gmail.com]
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
#######################################################################################

# This file implements an API endpoint for DIS background image removal system.
# [Self space] - [https://huggingface.co/spaces/leonelhs/removebg]
#
# Source code is based on or inspired by several projects.
# For more details and proper attribution, please refer to the following resources:
#
# - [DIS] - [https://github.com/xuebinqin/DIS]
# - [removebg] - [https://huggingface.co/spaces/gaviego/removebg]

import gradio as gr
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from huggingface_hub import hf_hub_download
from torch.autograd import Variable
from torchvision.transforms.functional import normalize
from itertools import islice

# project imports
from models.isnet import ISNetDIS

REPO_ID = "leonelhs/removators"

device = 'cuda' if torch.cuda.is_available() else 'cpu'

net = ISNetDIS()

model_path = hf_hub_download(repo_id=REPO_ID, filename='isnet.pth')
net.load_state_dict(torch.load(model_path, map_location=device))
net.to(device)
net.eval()

def im_preprocess(im, size):
    if len(im.shape) < 3:
        im = im[:, :, np.newaxis]
    if im.shape[2] == 1:
        im = np.repeat(im, 3, axis=2)
    if im.shape[2] == 4:
        im = im[:, :, :3]
    im_tensor = torch.tensor(im.copy(), dtype=torch.float32)
    im_tensor = torch.transpose(torch.transpose(im_tensor,1,2),0,1)
    if len(size)<2:
        return im_tensor, im.shape[0:2]
    else:
        im_tensor = torch.unsqueeze(im_tensor,0)
        im_tensor = F.interpolate(im_tensor, size, mode="bilinear")
        im_tensor = torch.squeeze(im_tensor,0)

    return im_tensor.type(torch.uint8), im.shape[0:2]


def predict(image):
    """
        Remove the background from an image.
        The function extracts the foreground and generates both a background-removed
        image and a binary mask.

        Parameters:
            image (string): File path to the input image.
        Returns:
            image (string): paths for image cutting mask.
    """
    im_tensor, shapes = im_preprocess(image, [1024, 1024])
    shapes = torch.from_numpy(np.array(shapes)).unsqueeze(0)

    im_tensor = torch.divide(im_tensor, 255.0)
    im_tensor = normalize(im_tensor, mean=[0.5, 0.5, 0.5], std=[1.0, 1.0, 1.0]).unsqueeze(0)
    im_tensor_v = Variable(im_tensor, requires_grad=False)  # wrap inputs in Variable
    ds_val = net(im_tensor_v)[0]  # list of 6 results
    prediction = ds_val[0][0, :, :, :]  # B x 1 x H x W    # we want the first one which is the most accurate prediction
    ## recover the prediction spatial size to the original image size
    size = (shapes[0][0], shapes[0][1])
    prediction = F.interpolate(torch.unsqueeze(prediction, 0), size, mode='bilinear')
    prediction = torch.squeeze(prediction)

    ma = torch.max(prediction)
    mi = torch.min(prediction)
    prediction = (prediction - mi) / (ma - mi)  # max = 1

    torch.cuda.empty_cache()
    return (prediction.detach().cpu().numpy() * 255).astype(np.uint8)  # it is the mask we need

def cuts(image):
    mask = predict(image)
    mask = Image.fromarray(mask).convert('L')
    cutted = Image.fromarray(image).convert("RGB")
    cutted.putalpha(mask)
    return [image, cutted], mask


with gr.Blocks(title="DIS") as app:
    navbar = gr.Navbar(visible=True, main_page_name="Workspace")
    gr.Markdown("## Dichotomous Image Segmentation")
    with gr.Row():
        with gr.Column(scale=1):
            inp_image = gr.Image(type="numpy", label="Upload Image")
            btn_predict = gr.Button(variant="primary", value="Remove background")
        with gr.Column(scale=2):
            with gr.Row():
                preview = gr.ImageSlider(type="filepath", label="Comparer")

    btn_predict.click(cuts, inputs=[inp_image], outputs=[preview, inp_image])

with app.route("Readme", "/readme"):
    with open("README.md") as f:
        for line in islice(f, 12, None):
            gr.Markdown(line.strip())


app.launch(share=False, debug=True, show_error=True, mcp_server=True, pwa=True)
app.queue()