app.py

import cv2
import gradio as gr
import os
from PIL import Image
import numpy as np
import torch
from torch.autograd import Variable
from torchvision import transforms
import torch.nn.functional as F
import gdown
import warnings
warnings.filterwarnings("ignore")

os.system("git clone https://github.com/xuebinqin/DIS")
os.system("mv DIS/IS-Net/* .")

# project imports
from data_loader_cache import normalize, im_reader, im_preprocess 
from models import *

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

# Download official weights
if not os.path.exists("saved_models"):
    os.mkdir("saved_models")
    os.system("mv isnet.pth saved_models/")

class GOSNormalize(object):
    '''
    Normalize the Image using torch.transforms
    '''
    def __init__(self, mean=[0.485,0.456,0.406], std=[0.229,0.224,0.225]):
        self.mean = mean
        self.std = std

    def __call__(self,image):
        image = normalize(image,self.mean,self.std)
        return image

transform = transforms.Compose([GOSNormalize([0.5,0.5,0.5],[1.0,1.0,1.0])])

def load_image(im_path, hypar):
    im = im_reader(im_path)
    im, im_shp = im_preprocess(im, hypar["cache_size"])
    im = torch.divide(im,255.0)
    shape = torch.from_numpy(np.array(im_shp))
    return transform(im).unsqueeze(0), shape.unsqueeze(0) # make a batch of image, shape

def build_model(hypar,device):
    net = hypar["model"]

    # convert to half precision
    if(hypar["model_digit"]=="half"):
        net.half()
        for layer in net.modules():
            if isinstance(layer, nn.BatchNorm2d):
                layer.float()

    net.to(device)

    if(hypar["restore_model"]!=""):
        net.load_state_dict(torch.load(hypar["model_path"]+"/"+hypar["restore_model"], map_location=device))
        net.to(device)
    net.eval()  
    return net

def predict(net, inputs_val, shapes_val, hypar, device):
    '''
    Given an Image, predict the mask
    '''
    net.eval()

    if(hypar["model_digit"]=="full"):
        inputs_val = inputs_val.type(torch.FloatTensor)
    else:
        inputs_val = inputs_val.type(torch.HalfTensor)

    inputs_val_v = Variable(inputs_val, requires_grad=False).to(device) # wrap inputs in Variable
    ds_val = net(inputs_val_v)[0] # list of 6 results
    pred_val = 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 orignal image size
    pred_val = torch.squeeze(F.upsample(torch.unsqueeze(pred_val,0),(shapes_val[0][0],shapes_val[0][1]),mode='bilinear'))

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

    if device == 'cuda': torch.cuda.empty_cache()
    return (pred_val.detach().cpu().numpy()*255).astype(np.uint8) # it is the mask we need
    
# Set Parameters
hypar = {} # parameters for inferencing

hypar["model_path"] ="./saved_models" ## load trained weights from this path
hypar["restore_model"] = "isnet.pth" ## name of the to-be-loaded weights
hypar["interm_sup"] = False ## indicate if activate intermediate feature supervision

##  choose floating point accuracy --
hypar["model_digit"] = "full" ## indicates "half" or "full" accuracy of float number
hypar["seed"] = 0

hypar["cache_size"] = [1024, 1024] ## cached input spatial resolution, can be configured into different size

## data augmentation parameters ---
hypar["input_size"] = [1024, 1024] ## model input spatial size, usually use the same value hypar["cache_size"], which means we don't further resize the images
hypar["crop_size"] = [1024, 1024] ## random crop size from the input, it is usually set as smaller than hypar["cache_size"], e.g., [920,920] for data augmentation

hypar["model"] = ISNetDIS()

# Build Model
net = build_model(hypar, device)

def inference(image):
    image_path = image
    image_tensor, orig_size = load_image(image_path, hypar) 
    mask = predict(net, image_tensor, orig_size, hypar, device)
  
    pil_mask = Image.fromarray(mask).convert('L')
    im_rgb = Image.open(image).convert("RGB")
  
    im_rgba = im_rgb.copy()
    im_rgba.putalpha(pil_mask)

    return [im_rgba, pil_mask]

# Translations for multi-language support
translations = {
    "pl": {
        "title": "Zaawansowane Segmentowanie Obrazów",
        "description": """
        # Zaawansowane Segmentowanie Obrazów

        **Zaawansowane Segmentowanie Obrazów** to zaawansowane narzędzie oparte na sztucznej inteligencji, zaprojektowane do precyzyjnego segmentowania obrazów. Aplikacja ta wykorzystuje najnowsze technologie głębokiego uczenia, aby generować dokładne maski dla różnych typów obrazów. Stworzona przez ekspertów, oferuje użytkownikom intuicyjny interfejs do przetwarzania obrazów. Niezależnie od tego, czy jest używana do celów zawodowych, czy do projektów osobistych, to narzędzie zapewnia najwyższą jakość i niezawodność w zadaniach segmentacji obrazów.
        """,
        "article": """
        ## Technologie

        - **Model**: ISNetDIS
        - **Stworzony przez**: Rafał Dembski
        - **Technologie**: PyTorch, Gradio, OpenCV
        """
    },
    "en": {
        "title": "Advanced Image Segmentation",
        "description": """
        # Advanced Image Segmentation

        **Advanced Image Segmentation** is a cutting-edge AI-powered tool developed to provide highly accurate image segmentation. This application utilizes state-of-the-art deep learning techniques to generate precise masks for various types of images. Developed by experts, it offers users an intuitive interface to process their images effortlessly. Whether for professional use or personal projects, this tool ensures superior quality and reliability in image segmentation tasks.
        """,
        "article": """
        ## Technologies

        - **Model**: ISNetDIS
        - **Developed by**: Rafał Dembski
        - **Technologies**: PyTorch, Gradio, OpenCV
        """
    },
    "de": {
        "title": "Fortschrittliche Bildsegmentierung",
        "description": """
        # Fortschrittliche Bildsegmentierung

        **Fortschrittliche Bildsegmentierung** ist ein fortschrittliches KI-gestütztes Tool, das entwickelt wurde, um hochpräzise Bildsegmentierung zu ermöglichen. Diese Anwendung nutzt modernste Techniken des Deep Learnings, um präzise Masken für verschiedene Arten von Bildern zu erstellen. Entwickelt von Experten, bietet es den Benutzern eine intuitive Oberfläche zur mühelosen Verarbeitung ihrer Bilder. Ob für professionelle Zwecke oder persönliche Projekte, dieses Tool gewährleistet höchste Qualität und Zuverlässigkeit bei Segmentierungsaufgaben.
        """,
        "article": """
        ## Technologien

        - **Modell**: ISNetDIS
        - **Entwickelt von**: Rafał Dembski
        - **Technologien**: PyTorch, Gradio, OpenCV
        """
    }
}

# Gradio setup with Monochrome theme, logo, and description with language support
css = """
#col-container {
    margin: 0 auto;
    max-width: 520px;
}
"""

def change_language(lang):
    return translations[lang]['title'], translations[lang]['description'], translations[lang]['article']

with gr.Blocks(theme=gr.themes.Monochrome(), css=css) as demo:
    language_selector = gr.Dropdown(choices=["pl", "en", "de"], value="en", label="Wybierz język / Select Language / Sprache auswählen", show_label=True)

    title = gr.Markdown(translations["en"]["title"])
    description = gr.Markdown(translations["en"]["description"])
    article = gr.Markdown(translations["en"]["article"])

    with gr.Row():
        prompt = gr.Image(type='filepath')
        result = gr.Image(label="Segmented Image", show_label=False)

    language_selector.change(
        fn=change_language,
        inputs=language_selector,
        outputs=[title, description, article],
    )

    gr.Interface(
        fn=inference,
        inputs=prompt,
        outputs=[result],
        title=title,
        description=description,
        article=article,
        allow_flagging='never',
        cache_examples=False,
    ).launch()