Update app.py

app.py

@@ -7,53 +7,51 @@ import torch
 from torch.autograd import Variable
 from torchvision import transforms
 import torch.nn.functional as F
+import gdown
 import warnings
+
 warnings.filterwarnings("ignore")
 
-# Sprawdzenie, czy katalog DIS istnieje przed klonowaniem
 if not os.path.exists("DIS"):
     os.system("git clone https://github.com/xuebinqin/DIS")
+if not os.path.exists("IS-Net"):
     os.system("mv DIS/IS-Net/* .")
 
 # project imports
-from data_loader_cache import normalize, im_reader, im_preprocess 
+from data_loader_cache import normalize, im_reader, im_preprocess
 from models import *
 
 # Helpers
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-# Sprawdzenie, czy katalog z modelami istnieje przed przeniesieniem
+# Download official weights
 if not os.path.exists("saved_models"):
     os.mkdir("saved_models")
 if not os.path.exists("saved_models/isnet.pth"):
     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]):
+    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)
+    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])])
+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)
+    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
+    return transform(im).unsqueeze(0), shape.unsqueeze(0)  # make a batch of image, shape
 
-def build_model(hypar,device):
+def build_model(hypar, device):
     net = hypar["model"]
 
-    # convert to half precision
-    if(hypar["model_digit"]=="half"):
+    if hypar["model_digit"] == "half":
         net.half()
         for layer in net.modules():
             if isinstance(layer, nn.BatchNorm2d):
@@ -61,53 +59,47 @@ def build_model(hypar,device):
 
     net.to(device)
 
-    if(hypar["restore_model"]!=""):
-        net.load_state_dict(torch.load(hypar["model_path"]+"/"+hypar["restore_model"], map_location=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()  
+    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"):
+    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
+    inputs_val_v = Variable(inputs_val, requires_grad=False).to(device)
+    ds_val = net(inputs_val_v)[0]
+    pred_val = ds_val[0][0, :, :, :]
 
-    # 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'))
+    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
+    pred_val = (pred_val - mi) / (ma - mi)
 
     if device == 'cuda': torch.cuda.empty_cache()
-    return (pred_val.detach().cpu().numpy()*255).astype(np.uint8) # it is the mask we need
-    
+    return (pred_val.detach().cpu().numpy() * 255).astype(np.uint8)
+
 # Set Parameters
-hypar = {} # parameters for inferencing
+hypar = {}
 
-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
+hypar["model_path"] = "./saved_models"
+hypar["restore_model"] = "isnet.pth"
+hypar["interm_sup"] = False
 
-##  choose floating point accuracy --
-hypar["model_digit"] = "full" ## indicates "half" or "full" accuracy of float number
+hypar["model_digit"] = "full"
 hypar["seed"] = 0
 
-hypar["cache_size"] = [1024, 1024] ## cached input spatial resolution, can be configured into different size
+hypar["cache_size"] = [1024, 1024]
 
-## 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["input_size"] = [1024, 1024]
+hypar["crop_size"] = [1024, 1024]
 
 hypar["model"] = ISNetDIS()
 
@@ -115,68 +107,57 @@ hypar["model"] = ISNetDIS()
 net = build_model(hypar, device)
 
 def inference(image):
-    image_path = image
-    image_tensor, orig_size = load_image(image_path, hypar) 
+    image_tensor, orig_size = load_image(image, 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
+# Translation texts
 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.
+        
+        **Technologie**:
+        - Model: ISNetDIS
+        - Stworzony przez: Rafał Dembski
+        - Technologie: PyTorch, Gradio, OpenCV
         """,
-        "article": """
-        ## Technologie
-
-        - **Model**: ISNetDIS
-        - **Stworzony przez**: Rafał Dembski
-        - **Technologie**: PyTorch, Gradio, OpenCV
-        """
+        "article": ""
     },
     "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.
+        **Advanced Image Segmentation** is an advanced AI-based tool designed for precise image segmentation. This application utilizes the latest deep learning technologies to generate accurate masks for different types of images. Created by experts, it offers users an intuitive interface for image processing. Whether used for professional purposes or personal projects, this tool ensures the highest quality and reliability in image segmentation tasks.
+        
+        **Technologies**:
+        - Model: ISNetDIS
+        - Created by: Rafał Dembski
+        - Technologies: PyTorch, Gradio, OpenCV
         """,
-        "article": """
-        ## Technologies
-
-        - **Model**: ISNetDIS
-        - **Developed by**: Rafał Dembski
-        - **Technologies**: PyTorch, Gradio, OpenCV
-        """
+        "article": ""
     },
     "de": {
-        "title": "Fortschrittliche Bildsegmentierung",
+        "title": "Fortgeschrittene 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.
+        **Fortgeschrittene Bildsegmentierung** ist ein fortschrittliches, auf künstlicher Intelligenz basierendes Werkzeug, das für die präzise Bildsegmentierung entwickelt wurde. Diese Anwendung nutzt die neuesten Technologien des Deep Learnings, um genaue Masken für verschiedene Bildtypen zu erzeugen. Von Experten erstellt, bietet es den Benutzern eine intuitive Benutzeroberfläche für die Bildverarbeitung. Ob für berufliche Zwecke oder persönliche Projekte, dieses Werkzeug gewährleistet höchste Qualität und Zuverlässigkeit bei der Bildsegmentierung.
+        
+        **Technologien**:
+        - Modell: ISNetDIS
+        - Erstellt von: Rafał Dembski
+        - Technologien: PyTorch, Gradio, OpenCV
         """,
-        "article": """
-        ## Technologien
-
-        - **Modell**: ISNetDIS
-        - **Entwickelt von**: Rafał Dembski
-        - **Technologien**: PyTorch, Gradio, OpenCV
-        """
+        "article": ""
     }
 }
 
-# Gradio setup with Monochrome theme, logo, and description with language support
 css = """
 #col-container {
     margin: 0 auto;
@@ -188,7 +169,10 @@ 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)
+    language = gr.State("en")
+
+    with gr.Row():
+        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"])