-
Notifications
You must be signed in to change notification settings - Fork 0
/
predict_part.py
139 lines (118 loc) · 4.44 KB
/
predict_part.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
import copy
import argparse
from PIL import Image
import torch
from detectron2.config import get_cfg
from detectron2.checkpoint import DetectionCheckpointer
from detectron2.data import MetadataCatalog
from detectron2.modeling import build_model
from detectron2.projects.deeplab import add_deeplab_config
from detectron2.utils.visualizer import Visualizer
from mask2former import add_maskformer2_config
import mask2former.data.transforms as T
class DefaultPredictorRGBA:
def __init__(self, cfg):
self.cfg = cfg.clone() # cfg can be modified by model
self.model = build_model(self.cfg)
self.model.eval()
if len(cfg.DATASETS.TEST):
self.metadata = MetadataCatalog.get(cfg.DATASETS.TEST[0])
checkpointer = DetectionCheckpointer(self.model)
checkpointer.load(cfg.MODEL.WEIGHTS)
self.aug = T.ResizeShortestEdgeRGBA(
[cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
)
self.input_format = cfg.INPUT.FORMAT
assert self.input_format == "RGBA", self.input_format
def __call__(self, original_image):
"""
Args:
original_image (np.ndarray): an image of shape (H, W, C) (in RGBA order).
Returns:
predictions (dict):
the output of the model for one image only.
See :doc:`/tutorials/models` for details about the format.
"""
with torch.no_grad(): # https://github.com/sphinx-doc/sphinx/issues/4258
height, width = original_image.shape[:2]
image = self.aug.get_transform(original_image).apply_image(original_image)
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
inputs = {"image": image, "height": height, "width": width}
predictions = self.model([inputs])[0]
return predictions
class InstancesWrap():
def __init__(self, instances, T=0.0, topk=100):
instances = copy.deepcopy(instances)
idx = instances.scores.argsort(descending=True)[:topk]
self.pred_scores = instances.scores[idx]
self.pred_masks = instances.pred_masks[idx, :, :]
idx = self.pred_scores > T
self.pred_scores = self.pred_scores[idx]
self.pred_masks = self.pred_masks[idx]
def has(self, name):
return hasattr(self, name)
def __len__(self):
return len(self.pred_scores)
def get_predictor(cfg_path, model_path):
cfg = get_cfg()
add_deeplab_config(cfg)
add_maskformer2_config(cfg)
cfg.merge_from_file(cfg_path)
cfg.MODEL.WEIGHTS = model_path
return DefaultPredictorRGBA(cfg)
def get_parser():
parser = argparse.ArgumentParser(description="predictor for Open Part Seg")
parser.add_argument(
"ckpt",
help="A path to checkpoint",
)
parser.add_argument(
"--input",
nargs="+",
help="A list of space separated input images; "
"or a single glob pattern such as 'directory/*.jpg'",
)
parser.add_argument(
"--output",
type=str,
default="demo/outputs",
help="A directory to save output visualizations",
)
parser.add_argument(
"--score-threshold",
type=float,
default=0.1,
help="Minimum score for instance predictions to be shown",
)
parser.add_argument(
"--topk",
type=int,
default=10,
help="Maximum number for instance predictions to be shown",
)
return parser
if __name__ == "__main__":
import os
import numpy as np
import warnings
args = get_parser().parse_args()
os.makedirs(args.output, exist_ok=True)
print(f"loading model from {args.ckpt}...")
predictor = get_predictor(
"configs/part_segmentation/dt_alpha/clsag.yaml",
args.ckpt,
)
for input in args.input:
im = np.asarray(Image.open(input))
assert im.shape[2] == 4, f"{input} needs to be RGBA image"
with warnings.catch_warnings():
warnings.simplefilter("ignore")
outputs = predictor(im)
instances = InstancesWrap(
outputs["instances"].to("cpu"), args.score_threshold, args.topk)
v = Visualizer(im[:, :, :3], scale=1.2)
out = v.draw_instance_predictions(instances)
out = Image.fromarray(out.get_image())
save_path = os.path.join(args.output, f"{os.path.basename(input)}")
print(f"saved to {save_path}")
out.save(save_path)