if __name__ == "__main__":
yaml_path = os.path.join(
Path(__file__).parent.parent, "params", "azure_centermask_SNU.yaml")
with open(yaml_path) as f:
params = yaml.load(f, Loader=yaml.FullLoader)
cfg = setup_cfg(params)
os.environ["CUDA_VISIBLE_DEVICES"] = params["is_gpu_id"]
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("==> Loading CenterMask on", device, params["is_gpu_id"])
demo = VisualizationDemo(cfg)
sock = su.initialize_client(params["tcp_ip"],
params["centermask_tcp_port"])
while True:
img = su.recvall_image(sock)
predictions, vis_output = demo.run_on_image(img)
pred_masks = predictions["instances"].pred_masks.cpu().detach().numpy(
) # (N, H, W),
pred_boxes = predictions["instances"].pred_boxes.tensor.cpu().detach(
).numpy() # (x1, y1, x2, y2)
pred_scores = predictions["instances"].scores.cpu().detach().numpy(
) # a vector of N confidence scores.
pred_classes = predictions["instances"].pred_classes.cpu().detach(
).numpy() # [0, num_categories).
vis_img = cv2.resize(vis_output.get_image()[:, :, ::-1],
(params["width"], params["height"]))
su.sendall_image(sock, vis_img)
su.sendall_pickle(sock, pred_masks)
su.sendall_pickle(sock, pred_boxes)
su.sendall_pickle(sock, pred_classes)
from easy_tcp_python2_3 import socket_utils as su
from imagenet_stubs.imagenet_2012_labels import label_to_name
if __name__ == "__main__":
sock = su.initialize_client('localhost', 7777)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = models.detection.maskrcnn_resnet50_fpn(pretrained=True).to(device)
model.eval()
print("Using mask R-CNN with ResNet-50 backbone", device)
transforms_composed = transforms.Compose([transforms.ToTensor()])
while True:
recv_image = su.recvall_image(sock)
x = transforms_composed(recv_image).unsqueeze(0)
outputs = model(x.to(device))[0]
probs = outputs['scores'].cpu().detach().numpy()
boxes = outputs["boxes"].cpu().detach().numpy()
labels = outputs["labels"].cpu().detach().numpy()
masks = outputs["masks"].cpu().detach().numpy()
keep = probs > 0.5
probs = probs[keep]
labels = labels[keep]
boxes = boxes[keep]
masks = masks[keep]
su.sendall_pickle(sock, probs)