alivegames/Grounded-Segment-Anything

1 2import warnings3warnings.filterwarnings('ignore')4 5import subprocess, io, os, sys, time6from loguru import logger7 8# os.system("pip install diffuser==0.6.0")9# os.system("pip install transformers==4.29.1")10 11os.environ["CUDA_VISIBLE_DEVICES"] = "0"12 13if os.environ.get('IS_MY_DEBUG') is None:14    result = subprocess.run(['pip', 'install', '-e', 'GroundingDINO'], check=True)15    print(f'pip install GroundingDINO = {result}')16 17# result = subprocess.run(['pip', 'list'], check=True)18# print(f'pip list = {result}')19 20sys.path.insert(0, './GroundingDINO')21 22import gradio as gr23 24import argparse25 26import copy27 28import numpy as np29import torch30from PIL import Image, ImageDraw, ImageFont, ImageOps31 32# Grounding DINO33import GroundingDINO.groundingdino.datasets.transforms as T34from GroundingDINO.groundingdino.models import build_model35from GroundingDINO.groundingdino.util import box_ops36from GroundingDINO.groundingdino.util.slconfig import SLConfig37from GroundingDINO.groundingdino.util.utils import clean_state_dict, get_phrases_from_posmap38 39import cv240import numpy as np41import matplotlib.pyplot as plt42from lama_cleaner.model_manager import ModelManager43from lama_cleaner.schema import Config as lama_Config44 45# segment anything46from segment_anything import build_sam, SamPredictor, SamAutomaticMaskGenerator47 48# diffusers49import PIL50import requests51import torch52from io import BytesIO53from diffusers import StableDiffusionInpaintPipeline54from huggingface_hub import hf_hub_download55 56from util_computer import computer_info57# relate anything58from ram_utils import iou, sort_and_deduplicate, relation_classes, MLP, show_anns, ram_show_mask59from ram_train_eval import RamModel,RamPredictor60from mmengine.config import Config as mmengine_Config61 62from app import *63 64config_file = 'GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py'65ckpt_repo_id = "ShilongLiu/GroundingDINO"66ckpt_filenmae = "groundingdino_swint_ogc.pth"67sam_checkpoint = './sam_vit_h_4b8939.pth' 68output_dir = "outputs"69device = 'cpu'70 71os.makedirs(output_dir, exist_ok=True)72groundingdino_model = None73sam_device = None74sam_model = None75sam_predictor = None76sam_mask_generator = None77sd_pipe = None78lama_cleaner_model= None79ram_model = None80kosmos_model = None81kosmos_processor = None82 83def get_args():84    argparser = argparse.ArgumentParser()85    argparser.add_argument("--input_image", "-i", type=str, default="", help="")86    argparser.add_argument("--text", "-t", type=str, default="", help="")87    argparser.add_argument("--output_image", "-o", type=str, default="", help="")88    argparser.add_argument("--cuda", "-c", type=str, default='cpu', help="cuda")89    args = argparser.parse_args()90    return args91 92# usage: 93#       python app_cli.py --input_image dog.png --text dog --output_image dog_remove.png94 95if __name__ == '__main__':96    args = get_args()97    logger.info(f'\nargs={args}\n')98 99    logger.info(f'loading models ... ')100    # set_device(args)  # If you have enough GPUs, you can open this comment101    load_groundingdino_model('cpu')102    load_sam_model(device)103    # load_sd_model(device)104    load_lama_cleaner_model(device)105    # load_ram_model(device)106 107    input_image = Image.open(args.input_image)108 109    run_rets = run_anything_task(input_image = input_image, 110                        text_prompt = args.text,  111                        task_type = 'remove', 112                        inpaint_prompt = '', 113                        box_threshold = 0.3, 114                        text_threshold = 0.25, 115                        iou_threshold = 0.8, 116                        inpaint_mode = "merge", 117                        mask_source_radio = "type what to detect below", 118                        remove_mode = "rectangle",   # ["segment", "rectangle"]119                        remove_mask_extend = "10", 120                        num_relation = 5,121                        kosmos_input = None,122                        cleaner_size_limit = -1,123                        )124    output_images = run_rets[0]125    if len(output_images) > 0:126        logger.info(f'save result to {args.output_image} ... ')        127        output_images[-1].save(args.output_image)128        # count = 0129        # for output_image in output_images:130        #     count += 1131        #     if isinstance(output_image, np.ndarray):132        #         output_image = PIL.Image.fromarray(output_image.astype(np.uint8))133        #     output_image.save(args.output_image.replace(".",  f"_{count}."))134