Spaces:
Sleeping
Sleeping
import matplotlib.pyplot as plt | |
import matplotlib.pylab as pylab | |
import requests | |
from io import BytesIO | |
from PIL import Image | |
import numpy as np | |
pylab.rcParams['figure.figsize'] = 20, 12 | |
from maskrcnn_benchmark.config import cfg | |
from maskrcnn_benchmark.engine.predictor_glip import GLIPDemo | |
import argparse | |
import pdb | |
def load(url_or_path): | |
""" | |
Given an url or a path, this loads the file and | |
""" | |
if url_or_path.startswith("http"): | |
response = requests.get(url_or_path) | |
pil_image = Image.open(BytesIO(response.content)).convert("RGB") | |
# convert to BGR format | |
image = np.array(pil_image)[:, :, [2, 1, 0]] | |
else: | |
image = np.array(Image.open(url_or_path).convert("RGB"))[:, :, [2, 1, 0]] | |
return image | |
parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference") | |
parser.add_argument("--config", default="configs/pretrain/glip_Swin_T_O365_GoldG.yaml", metavar="FILE", help="path to config file", type=str) | |
parser.add_argument("--weight", default="OUTPUTS/GLIP_MODEL4/model_0020000.pth", metavar="FILE", help="path to weight file", type=str) | |
parser.add_argument("--image", default="http://farm4.staticflickr.com/3693/9472793441_b7822c00de_z.jpg", metavar="FILE", help="path to weight file", type=str) | |
parser.add_argument("--conf", default=0.4, type=float) | |
parser.add_argument("--caption", default="", type=str) | |
parser.add_argument("--ground_tokens", default=None, type=str) | |
args = parser.parse_args() | |
# update the config options with the config file | |
# manual override some options | |
cfg.local_rank = 0 | |
cfg.num_gpus = 1 | |
cfg.merge_from_file(args.config) | |
cfg.merge_from_list(["MODEL.WEIGHT", args.weight]) | |
cfg.merge_from_list(["MODEL.DEVICE", "cuda"]) | |
glip_demo = GLIPDemo( | |
cfg, | |
min_image_size=800, | |
show_mask_heatmaps=False | |
) | |
athetics_params = { | |
"skip_name": False, # whether we overlay the phrase over the box | |
"override_color": (255, 255, 255), # box color, default is white | |
"text_size": 1.0, | |
"text_pixel": 3, | |
"box_alpha": 1.0, | |
"box_pixel": 5, | |
"text_offset_original": 8, # distance between text and box | |
} | |
image = load(args.image) | |
specified_tokens = args.ground_tokens.split(";") if args.ground_tokens is not None else None | |
result, _ = glip_demo.run_on_web_image( | |
image, | |
args.caption, | |
args.conf, | |
specified_tokens, | |
**athetics_params) | |
plt.imshow(result[:, :, [2, 1, 0]]) | |
plt.axis("off") | |
plt.savefig(args.image.replace('.png', "_demo.png").replace('.jpg', "_demo.jpg").replace('.jpeg', "_demo.jpeg"), bbox_inches='tight', pad_inches=0) # save as xxx_demo.xxx | |