207 lines
8.3 KiB
Python
207 lines
8.3 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
import base64, os
|
|
import io
|
|
import random
|
|
from PIL import Image
|
|
import torch
|
|
import logging
|
|
import easyocr
|
|
|
|
import folder_paths
|
|
from .utils import get_config,load_list_images,convert_safetensor_to_pt,tensor2imglist
|
|
from .OmniParser.utils import check_ocr_box, get_yolo_model, get_caption_model_processor, get_som_labeled_img
|
|
|
|
node_cur_path = os.path.dirname(os.path.abspath(__file__))
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
class OmniParser_Loader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"repo": ("STRING", {"default": "microsoft/OmniParser", }),
|
|
"platform": (["pc", "web","mobile"],),
|
|
},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("OP_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "main"
|
|
CATEGORY = "OmniParser"
|
|
|
|
def main(self, repo,platform,):
|
|
if repo:
|
|
if "florence" in repo.lower():
|
|
model_name="florence2"
|
|
elif "blip" in repo.lower():
|
|
model_name="blip2"
|
|
elif repo== "microsoft/OmniParser":
|
|
model_name = "default"
|
|
else:
|
|
raise "unsupport repo id"
|
|
else:
|
|
raise "must fill repo or local path"
|
|
|
|
yolo_model_converted=os.path.join(node_cur_path,"OmniParser/weights/icon_detect/best.pt")
|
|
if not os.path.exists(yolo_model_converted):
|
|
convert_safetensor_to_pt()
|
|
logging.info("convert_safetensor_to_pt is done.")
|
|
|
|
yolo_model = get_yolo_model(model_path=yolo_model_converted)
|
|
|
|
caption_model_processor = get_caption_model_processor(model_name=model_name,
|
|
model_name_or_path=repo)
|
|
draw_bbox_config=get_config(platform)
|
|
|
|
logging.info("loading checkpoint done.")
|
|
model={"yolo_model":yolo_model,"caption_model_processor":caption_model_processor,"draw_bbox_config":draw_bbox_config}
|
|
return (model,)
|
|
|
|
class OmniParser_Sampler:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model":("OP_MODEL",),
|
|
"image": ("IMAGE",), # [B,H,W,C], C=3
|
|
"text_threshold":("FLOAT", {
|
|
"default": 0.9,
|
|
"min": 0.1,
|
|
"max": 1.0,
|
|
"step": 0.1,
|
|
"display": "number",}),
|
|
"box_threshold": ("FLOAT", {
|
|
"default": 0.05,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number", }),
|
|
"iou_threshold": ("FLOAT", {
|
|
"default": 0.1,
|
|
"min": 0.01,
|
|
"max": 1.0,
|
|
"step": 0.01,
|
|
"display": "number", }),
|
|
"orc_model": (["easyocr", "paddleocr"],),
|
|
"orc_mode": (["en","ch_sim","ch_tra","ja","ko","ta","te","kn","th","ru","ar","hi"],),
|
|
"output_bb_format": (["xyxy", "xywh"],),
|
|
"display_img": ("BOOLEAN", {"default": False},),
|
|
"paragraph": ("BOOLEAN", {"default": False},),
|
|
"output_coord_in_ratio": ("BOOLEAN", {"default": True},),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE","STRING")
|
|
RETURN_NAMES = ("image","string")
|
|
FUNCTION = "main"
|
|
CATEGORY = "OmniParser"
|
|
|
|
def main(self, model,image,text_threshold,box_threshold,iou_threshold,orc_model,orc_mode,output_bb_format,display_img,paragraph,output_coord_in_ratio,**kwargs):
|
|
#optional chocie
|
|
optional_c=True
|
|
|
|
#pre model
|
|
yolo_model=model.get("yolo_model")
|
|
draw_bbox_config=model.get("draw_bbox_config")
|
|
caption_model_processor=model.get("caption_model_processor")
|
|
|
|
if orc_model=="easyocr":
|
|
ORC = easyocr.Reader([orc_mode])
|
|
use_paddleocr = False
|
|
else:
|
|
if orc_mode=="ch_sim":
|
|
orc_mode="ch"
|
|
elif orc_mode=="ch_tra":
|
|
orc_mode = "chinese_cht"
|
|
elif orc_mode=="ko":
|
|
orc_mode = "korean"
|
|
elif orc_mode=="ja":
|
|
orc_mode = "japan"
|
|
elif orc_mode=="ru":
|
|
orc_mode = "cyrillic"
|
|
elif orc_mode == "ar":
|
|
orc_mode = " arabic"
|
|
elif orc_mode == "en":
|
|
orc_mode = "en"
|
|
elif orc_mode == "hi":
|
|
orc_mode = "devanagari"
|
|
else:
|
|
orc_mode = "latin"
|
|
|
|
use_paddleocr = True #another choice
|
|
from paddleocr import PaddleOCR
|
|
ORC = PaddleOCR(
|
|
lang=orc_mode, # other lang also available
|
|
use_angle_cls=False,
|
|
use_gpu=False, # using cuda will conflict with pytorch in the same process
|
|
show_log=False,
|
|
max_batch_size=1024,
|
|
use_dilation=True, # improves accuracy
|
|
det_db_score_mode='slow', # improves accuracy
|
|
rec_batch_num=1024)
|
|
|
|
#pre image if list
|
|
file_prefix = ''.join(random.choice("0123456789abcdefg") for _ in range(6))
|
|
image_list,B = tensor2imglist(image, np_out=True)
|
|
path_list = []
|
|
if B==1:
|
|
image_input = image_list[0]
|
|
image_save_path = os.path.join(folder_paths.get_input_directory(), f'Temp_{file_prefix}.png')
|
|
image_input.save(image_save_path)
|
|
path_list=[image_save_path]
|
|
else:
|
|
for i,img in enumerate(image_list):
|
|
image_save_path = os.path.join(folder_paths.get_input_directory(), f'Temp_{file_prefix}_{i}.png')
|
|
img.save(image_save_path)
|
|
path_list.append(image_save_path)
|
|
|
|
ouput_img=[]
|
|
ouput_text=[]
|
|
for batch_num,image_path in enumerate(path_list):
|
|
# pre box
|
|
ocr_bbox_rslt, is_goal_filtered = check_ocr_box(image_path, ORC, use_paddleocr,
|
|
display_img=display_img, output_bb_format=output_bb_format,
|
|
goal_filtering=None,
|
|
easyocr_args={'paragraph': paragraph,
|
|
'text_threshold': text_threshold})
|
|
text, ocr_bbox = ocr_bbox_rslt
|
|
# infer
|
|
logging.info("start infer..")
|
|
dino_labled_img, label_coordinates, parsed_content_list = get_som_labeled_img(optional_c, image_path,
|
|
yolo_model,batch_num,
|
|
BOX_TRESHOLD=box_threshold,
|
|
output_coord_in_ratio=output_coord_in_ratio,
|
|
ocr_bbox=ocr_bbox,
|
|
draw_bbox_config=draw_bbox_config,
|
|
caption_model_processor=caption_model_processor,
|
|
ocr_text=text,
|
|
iou_threshold=iou_threshold)
|
|
image = Image.open(io.BytesIO(base64.b64decode(dino_labled_img)))
|
|
ouput_img.append(image)
|
|
content_list = '\n'.join(parsed_content_list)
|
|
ouput_text.append(content_list)
|
|
print('finish processing')
|
|
image=load_list_images(ouput_img)
|
|
prompt='\n\n'.join(ouput_text)
|
|
return (image,prompt)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"OmniParser_Loader": OmniParser_Loader,
|
|
"OmniParser_Sampler":OmniParser_Sampler,
|
|
}
|
|
|
|
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"OmniParser_Loader": "OmniParser_Loader",
|
|
"OmniParser_Sampler":"OmniParser_Sampler",
|
|
}
|