Skip to content

Commit

Permalink
Fix duplicate tabs on WebUI startup
Browse files Browse the repository at this point in the history
Fixes issue mattyamonaca#36 caused by double running code due to import.
mattyamonaca#36
  • Loading branch information
predprey2 committed May 23, 2023
1 parent b179be6 commit 7367898
Show file tree
Hide file tree
Showing 3 changed files with 26 additions and 22 deletions.
2 changes: 1 addition & 1 deletion scripts/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@
from pydantic import BaseModel, Field

from modules.api.api import encode_pil_to_base64, decode_base64_to_image
from scripts.main import process_image, model_list
from scripts.utils import process_image, model_list


class PBRRemRequest(BaseModel):
Expand Down
22 changes: 1 addition & 21 deletions scripts/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@
import io
import json
import numpy as np
import cv2
import re

import gradio as gr
Expand All @@ -12,24 +11,11 @@
import modules.shared as shared
from modules import images

from scripts.td_abg import get_foreground
from scripts.convertor import pil2cv
from scripts.batch_dir import save_image_dir, modify_basename
from scripts.utils import process_image, model_list

try:
from modules.paths_internal import extensions_dir
except Exception:
from modules.extensions import extensions_dir

from collections import OrderedDict
from PIL import Image

model_cache = OrderedDict()
sam_model_dir = os.path.join(
extensions_dir, "PBRemTools/models/")
model_list = [f for f in os.listdir(sam_model_dir) if os.path.isfile(
os.path.join(sam_model_dir, f)) and f.split('.')[-1] != 'txt']

def processing(single_image, batch_image, input_dir, output_dir, output_mask_dir, show_result, input_tab_state, *rem_args):
# 0: single
if (input_tab_state == 0):
Expand Down Expand Up @@ -78,12 +64,6 @@ def processing(single_image, batch_image, input_dir, output_dir, output_mask_dir
else:
return None

def process_image(target_image, *rem_args):
image = pil2cv(target_image)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
mask, image = get_foreground(image, *rem_args)
return image, mask

class Script(scripts.Script):
def __init__(self) -> None:
super().__init__()
Expand Down
24 changes: 24 additions & 0 deletions scripts/utils.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
import os
import cv2

from scripts.td_abg import get_foreground
from scripts.convertor import pil2cv

try:
from modules.paths_internal import extensions_dir
except Exception:
from modules.extensions import extensions_dir

from collections import OrderedDict

model_cache = OrderedDict()
sam_model_dir = os.path.join(
extensions_dir, "PBRemTools/models/")
model_list = [f for f in os.listdir(sam_model_dir) if os.path.isfile(
os.path.join(sam_model_dir, f)) and f.split('.')[-1] != 'txt']

def process_image(target_image, *rem_args):
image = pil2cv(target_image)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
mask, image = get_foreground(image, *rem_args)
return image, mask

0 comments on commit 7367898

Please sign in to comment.