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