def generalInference(model, im_names, predictor):
# num_devices = get_world_size()
with inference_context(model), torch.no_grad():
for im_name in tqdm(im_names):
im = Image.open(im_name)
w, h = im.size
# imSize = (h, w)
# mode = im.mode
im_data = np.array(im)
# Convert RGB to BGR
open_cv_image = im_data[:, :, ::-1].copy()
xml_name = os.path.join(
save_dir,
os.path.splitext(osp.basename(im_name))[0] + ".xml")
if os.path.exists(xml_name):
print("exists:", xml_name)
else:
output = predictor(open_cv_image)
v = Visualizer(im_data, scale=0.5)
out = v.draw_instance_predictions(
output["instances"].to("cpu"))
cv2.imshow("show_instances", out.get_image()[:, :, ::-1])
if "instances" in output:
instances = output["instances"] #.to(self.cpu_device)
prediction = instances.get_fields()
boxes = prediction['pred_boxes']
scores = prediction['scores']
classes = prediction['pred_classes']
# select_top() 选出top预测结果
boxes = boxes.tensor.cpu().numpy()
scores = scores.cpu().numpy()
classes = classes.cpu().numpy()
# wbf() weighted boxes fusion调用
saveGaofenXMLFormat(xml_name, im_name, boxes, scores)
def do_prediction_and_visualization(dataset, cfg):
dataset_dicts = DatasetCatalog.get(dataset)
metadata = MetadataCatalog.get(dataset)
# label=MetadataCatalog.get('cuboid_dataset_val').thing_classes
predictor = get_predictor(dataset, cfg)
# for d in random.sample(dataset_dicts, 50):
for d in dataset_dicts:
# if d['annotations']: #display +ve images only
im = cv2.imread(d["file_name"])
outputs = predictor(im)
# print(outputs)
v = Visualizer(
im[:, :, ::-1],
metadata=metadata,
scale=1,
)
output_instances = outputs['instances'].to('cpu')
pred = output_instances.pred_classes
# print(pred.tolist())
classes = []
for i in range(len(pred.tolist())):
classes.append('cuboid')
scores = output_instances.scores
labels = [
"{} {:.0f}%".format(l, s * 100) for l, s in zip(classes, scores)
]
# gt = v.draw_dataset_dict(d) #display ground truth annotation
out = v.overlay_instances(boxes=output_instances.pred_boxes,
labels=labels,
keypoints=output_instances.pred_keypoints)
final_img = cv2.resize(out.get_image()[:, :, ::-1], (900, 900))
cv2.imshow('Predication: ' + d['image_id'] + '.jpg', final_img)
k = cv2.waitKey(0)
if k == 27: #esc key for stop
cv2.destroyAllWindows()
break
cv2.destroyAllWindows()
def detect_items_in_image(img, predictor, item_metadata):
target_items_dict = {}
s_time = time.time()
outputs = predictor(img)
instances = outputs["instances"].to("cpu")
# format is documented at
# https://detectron2.readthedocs.io/tutorials/models.html # model-output-format
v = Visualizer(
img[:, :, ::-1],
metadata=item_metadata,
scale=0.5,
instance_mode=ColorMode.IMAGE_BW
# remove the colors of unsegmented pixels.
# This option is only available for segmentation models
)
out = v.draw_instance_predictions(instances)
masks_arr = instances.pred_masks.numpy()
item_classes = instances.pred_classes.numpy()
e_time = time.time()
print('detectron2 using : ',
int(round(e_time * 1000)) - int(round(s_time * 1000)))
if item_classes.any():
for i, item_num in enumerate(item_classes):
print(CLASS_NAMES[item_num])
arr_h, arr_w = np.nonzero(masks_arr[i])
mask_points = np.array([arr_h, arr_w]).T
print(f"mask_points:{mask_points.shape}")
target_items_dict[CLASS_NAMES[item_num]] = [mask_points]
res_img = out.get_image()
plt.imshow(res_img)
plt.savefig("result/0.jpg")
plt.show()
return target_items_dict
def checkout_dataset_annotation(save_path="/root/data/gvision/dataset/train_all_annos/s0.3_t0.7_all/pv_image_in_annos",subdata_name="pv_train",showwidth=640):
"""
Create configs and perform basic setups.
DATASETS:person(vehicle or pv)_VAL(TRAIN or train) or 查看融合的结果
"""
os.makedirs(save_path,exist_ok=True)
metadata_dicts= MetadataCatalog.get(subdata_name)
print("checkout_dataset_annotation------------------------------------------------start")
dataset_dicts = DatasetCatalog.get(subdata_name)
# random.seed(0)
# for d in dataset_dicts:
for d in random.sample(dataset_dicts, 10):
img = cv2.imread(d["file_name"])
visualizer = Visualizer(img[:, :, ::-1], metadata=metadata_dicts, scale=1)#font sacle
vis = visualizer.draw_dataset_dict(d)
save_name=os.path.join(save_path,"visual_{}".format(os.path.basename(d["file_name"])))
print(save_name)
img=vis.get_image()[:, :, ::-1]
imgheight, imgwidth = img.shape[:2]
scale = showwidth / imgwidth
img = cv2.resize(vis.get_image()[:, :, ::-1], (int(imgwidth * scale), int(imgheight * scale)),interpolation=cv2.INTER_AREA)
cv2.imwrite(save_name,img ,[int(cv2.IMWRITE_PNG_COMPRESSION), 9])
print("checkout_dataset{}_annotation------------------------------------------------end".format(subdata_name))
"""
def inference(self, file):
predictor = DefaultPredictor(self.cfg)
im = cv.imread(file)
outputs = predictor(im)
# with open(self.curr_dir+'/data.txt', 'w') as fp:
# json.dump(outputs['instances'], fp)
# # json.dump(cfg.dump(), fp)
# get metadata
metadata = MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0])
# visualise
v = Visualizer(im[:, :, ::-1], metadata=metadata, scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
# get image
img = Image.fromarray(np.uint8(v.get_image()[:, :, ::-1]))
# write to jpg
# cv.imwrite('img.jpg',v.get_image())
return img
def create_dataset_files():
img_dicts_train = get_image_dicts(image_folder='dataset/train/',
annotation_file="dataset/" +
config.ANNOTATION_TRAIN_CSV,
target_classes=["Parrot"])
img_dicts_validate = get_image_dicts(image_folder='dataset/validate/',
annotation_file="dataset/" +
config.ANNOTATION_VALIDATE_CSV,
target_classes=["Parrot"])
loaded_train_img_dicts = load_json_labels('dataset/train')
register_parrot_dataset()
# Setup metadata variable
parrot_metadata = MetadataCatalog.get("parrot_train")
for label in random.sample(img_dicts_train, 3):
img = cv2.imread(label["file_name"])
visualizer = Visualizer(
img[:, :, ::-1], # revervse the image pixel order (BGR -> RGB)
metadata=
parrot_metadata, # use the metadata variable we created about
scale=0.5)
vis = visualizer.draw_dataset_dict(label)
cv2.imshow('predicted', vis.get_image()[:, :, ::-1])
cv2.waitKey(0)
def plot_segmentation_masks_beauty(result, image):
# We extract the segments info and the panoptic result from DETR's prediction
segments_info = deepcopy(result["segments_info"])
# Panoptic predictions are stored in a special format png
panoptic_seg = Image.open(io.BytesIO(result['png_string']))
final_w, final_h = panoptic_seg.size
# We convert the png into an segment id map
panoptic_seg = np.array(panoptic_seg, dtype=np.uint8)
panoptic_seg = torch.from_numpy(rgb2id(panoptic_seg))
# Detectron2 uses a different numbering of coco classes, here we convert the class ids accordingly
meta = MetadataCatalog.get("coco_2017_val_panoptic_separated")
for i in range(len(segments_info)):
c = segments_info[i]["category_id"]
segments_info[i]["category_id"] = meta.thing_dataset_id_to_contiguous_id[c] if segments_info[i]["isthing"] else meta.stuff_dataset_id_to_contiguous_id[c]
# Finally we visualize the prediction
v = Visualizer(np.array(image.copy().resize((final_w, final_h)))[:, :, ::-1], meta, scale=1.0)
v._default_font_size = 20
v = v.draw_panoptic_seg_predictions(panoptic_seg, segments_info, area_threshold=0)
return v.get_image()
# ----------------------------------------------------------------------------------
def get_keypoints(self, frame):
outputs = self.predictor(frame)
v = Visualizer(frame,
MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0]),
scale=1)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
instances = outputs["instances"].to("cpu")
bbox = []
keypoints = []
pred_boxes = instances._fields["pred_boxes"].tensor.numpy()
pred_keypoints = instances._fields["pred_keypoints"].numpy()
for i in range(len(instances._fields["scores"])):
if instances._fields["scores"][i] > 0.99 and pred_keypoints[
i].shape[0] == 17:
bbox.append(pred_boxes[i])
keypoints.append(pred_keypoints[i])
keypoints = np.array(keypoints)
bbox = np.array(bbox)
eyes = Detectron2Keypoints.keypoints_to_eyes(
keypoints) if len(keypoints) > 0 else []
head_bbox = Detectron2Keypoints.bbox_keypoints_head_bbox(
bbox, keypoints) if len(keypoints) > 0 else []
return out.get_image(), head_bbox, eyes
def predictDataset(datasetName, dataset_data, dataset_metadata):
"""
:param datasetName: name of the registered dataset
:param dataset_data: DatasetCatalog of the registered dataset
:param dataset_metadata: MetadatasetCatalog of the registerd data
:return: shows 10 randomly selected pictures using opencv2. the ungemented portions are greyed out.
press 0 key to see next image
"""
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.DATALOADER.NUM_WORKERS = 2
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.02
cfg.SOLVER.MAX_ITER = 300
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 64
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set the testing threshold for this model
cfg.DATASETS.TEST = (datasetName, )
predictor = DefaultPredictor(cfg)
for d in random.sample(dataset_data, 10):
im = cv2.imread(d["file_name"])
outputs = predictor(im)
v = Visualizer(
im[:, :, ::-1],
metadata=dataset_metadata,
scale=0.8,
instance_mode=ColorMode.
IMAGE_BW # remove the colors of unsegmented pixels
)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imshow("image", v.get_image()[:, :, ::-1])
cv2.waitKey(0)
def visualize_element_prediction(img,
metadata,
predictions,
shown_categories,
score_cutoff=0):
vis = Visualizer(img, metadata)
# get targeted annotations to show
boxes = []
labels = []
# get the specific categories to show
for idx in range(len(predictions)):
if float(predictions.iloc[idx]
["score"]) >= score_cutoff and predictions.iloc[idx][
"category_id"] in shown_categories:
boxes.append(predictions.iloc[idx]["bbox"])
labels.append(predictions.iloc[idx]["category_id"])
names = metadata.get("thing_classes", None)
if names:
labels = [names[i] for i in labels]
boxes = np.array(boxes, np.float).reshape((-1, 4))
vis_gt = vis.overlay_instances(labels=labels, boxes=boxes).get_image()
del boxes, labels
return vis_gt[:, :, ::-1]
def image_predict(imgfile):
from PIL import Image
import numpy
from detectron2.utils.visualizer import ColorMode
import time
image = Image.open(imgfile)
bgr_img = cv2.cvtColor(numpy.asarray(image), cv2.COLOR_RGB2BGR)
num_class = len(class_list)
# MetadataCatalog.get(regist_train_name).thing_classes = class_list
MetadataCatalog.get(regist_train_name)
train_metadata = MetadataCatalog.get(regist_train_name)
start_time = time.time()
print('预测开始')
outputs = PREDICTOR(bgr_img)
print('预测完成', time.time() - start_time)
pred_scores = outputs["instances"].scores.tolist()
pred_classes = outputs["instances"].pred_classes.data.cpu().numpy().tolist(
) if outputs["instances"].get_fields()["pred_boxes"] else None
pred_bboxes = outputs["instances"].pred_boxes.tensor.tolist(
) if outputs["instances"].get_fields()["pred_boxes"] else None
v = Visualizer(bgr_img[:, :, ::-1],
metadata=train_metadata,
scale=1,
instance_mode=ColorMode.IMAGE_BW
) # remove the colors of unsegmented pixels)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
print(pred_classes)
savr_dir = "./test_out2/"
save_file = os.path.join(savr_dir, os.path.basename(imgfile) + "_pre.png")
cv2.imwrite(save_file, v.get_image()[:, :, ::-1])
def predict(photo_name):
"""
photo_name - Сылка на фото.
return - bio -> фото с отмеченными линиями
return - Аннотации к линиям
"""
im = cv2.imread(photo_name)
outputs = predictor(im)
v = Visualizer(
im[:, :, ::-1],
metadata=clothesnet,
scale=0.8,
)
outputs["instances"].remove('pred_boxes')
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
is_success, buffer = cv2.imencode(".jpg", v.get_image()[:, :, ::-1])
bio = io.BytesIO(buffer)
bio.name = 'image.jpeg'
bio.seek(0)
text = sorted(set(np.array(outputs["instances"].pred_classes)),
reverse=True)
return bio, boxes(text)
def segmentate(image_path, mode='box'):
image = cv2.imread(image_path)
outputs = predictor(image)
# return None if no fashion objects found on picture
if not outputs['instances'].pred_boxes:
return None
v = Visualizer(image[:, :, ::-1], metadata=fashion_metadata)
vis_full = v.draw_instance_predictions(outputs['instances'].to('cpu'))
result_images_dir = Path(FILES_DIR, str(uuid.uuid4()))
result_images_dir.mkdir(exist_ok=True, parents=True)
full_image_path = Path(result_images_dir,
f"full_{uuid.uuid4()}.jpg").as_posix()
cv2.imwrite(full_image_path, vis_full.get_image()[:, :, ::-1])
if mode == 'box':
crop_image_and_save(image=image,
outputs=outputs,
result_images_dir=result_images_dir)
elif mode == 'mask':
crop_image_mask_and_save(image=image,
outputs=outputs,
result_images_dir=result_images_dir,
vis=v)
# make archive to send many pictures
shutil.make_archive(
result_images_dir,
'zip',
root_dir=Path(FILES_DIR, result_images_dir),
)
shutil.rmtree(Path(FILES_DIR, result_images_dir))
return result_images_dir
def process_single(metadata, dic, args, queue:Queue=None):
masks = BorderMasks([x['segmentation'] for x in dic['annotations']
if not isinstance(x['segmentation'], dict)])
img = utils.read_image(dic["file_name"], "RGB")
borders, centers, sizes = masks.masks(mask_size=img.shape[:2])
if args.show:
visualizer = Visualizer(img, metadata=metadata)
vis = visualizer.draw_dataset_dict(dic)
webcv2.imshow(dic["file_name"], vis.get_image()[:, :, ::-1])
webcv2.imshow(dic["file_name"] + "-border", borders * 255)
webcv2.imshow(dic["file_name"] + "-centers", centers * 255)
webcv2.imshow(dic["file_name"] + "-sizes", (sizes / sizes.max()).sum(-1) * 255)
webcv2.waitKey()
else:
file_name = os.path.basename(dic["file_name"])
if queue:
queue.put(("borders", borders.tostring(), file_name))
queue.put(("centers", centers.tostring(), file_name))
queue.put(("sizes", sizes.tostring(), file_name))
else:
save(borders, cfg.MASK_DIRECTORY, split_name, "borders", file_name)
save(centers, cfg.MASK_DIRECTORY, split_name, "centers", file_name)
save(sizes, cfg.MASK_DIRECTORY, split_name, "sizes", file_name)
def inference(im, predictor, data_metadata, threshold):
# Convert PIL image to array
im = np.asarray(im)
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1], metadata=data_metadata, scale=0.5)
# take only predictions with 25% confidence only for RetinaNet as they tend to overdraw
filtered = outputs['instances'].to("cpu")._fields
filtered_idx = []
for i in range(len(filtered['scores'])):
if filtered['scores'][i] >= threshold:
filtered_idx.append(i)
filt_instance = Instances(
image_size=(im.shape[0], im.shape[1]),
pred_boxes=outputs['instances']._fields['pred_boxes'][filtered_idx],
pred_classes=outputs['instances']._fields['pred_classes']
[filtered_idx],
scores=outputs['instances']._fields['scores'][filtered_idx])
v = v.draw_instance_predictions(filt_instance.to("cpu"))
return v.get_image(), filt_instance
def run_on_image(self, image):
"""
Args:
image (np.ndarray): an image of shape (H, W, C) (in BGR order).
This is the format used by OpenCV.
Returns:
predictions (dict): the output of the model.
vis_output (VisImage): the visualized image output.
"""
vis_output = None
predictions = self.predictor(image)
# Convert image from OpenCV BGR format to Matplotlib RGB format.
image = image[:, :, ::-1]
visualizer = Visualizer(image,
self.metadata,
instance_mode=self.instance_mode)
if "instances" in predictions:
instances = predictions["instances"].to(self.cpu_device)
vis_output = visualizer.draw_instance_predictions(
predictions=instances)
return predictions, vis_output
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
outputs = predictor(im)
# look at the outputs. See https://detectron2.readthedocs.io/tutorials/models.html#model-output-format for specification
print(outputs["instances"].pred_classes)
print(outputs["instances"].pred_boxes)
# We can use `Visualizer` to draw the predictions on the image.
v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
plt.imshow(out.get_image()[:, :, ::-1])
# plt.show()
'''
dataset
'''
from detectron2.structures import BoxMode
def get_balloon_dicts(img_dir):
json_file = os.path.join(img_dir, "via_region_data.json")
with open(json_file) as f:
imgs_anns = json.load(f)
cfg.DATASETS.TEST = ("cats_val", )
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.7
predictor = DefaultPredictor(cfg)
dataset_dicts = get_data_dicts("../data/validation")
for idx, d in enumerate(random.sample(dataset_dicts, 3)):
im = cv2.imread(d["file_name"])
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
metadata=MetadataCatalog.get("cats_val").set(
thing_classes=classes,
thing_colors=[(177, 205, 223), (223, 205, 177)]),
scale=0.8,
instance_mode=ColorMode.IMAGE_BW)
pred_class = (outputs['instances'].pred_classes).detach()[0]
pred_score = (outputs['instances'].scores).detach()[0]
print(f"File: {d['file_name']}")
print(f"--> Class: {classes[pred_class]}, {pred_score * 100:.2f}%")
# Save image predictions
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
image_name = f"inference_{idx}.jpg"
cv2.imwrite(image_name, v.get_image()[:, :, ::-1])
drive.mount('/content/drive')
# Commented out IPython magic to ensure Python compatibility.
# %cd /content/drive/My Drive/STAT6000/ResearchProject/train2
#if your dataset is in COCO format, this cell can be replaced by the following three lines:
from detectron2.data.datasets import register_coco_instances
register_coco_instances("pods_train", {}, "instances_default.json", "/content/drive/My Drive/STAT6000/ResearchProject/train2")
register_coco_instances("pods_test", {}, "instances_default2.json", "/content/drive/My Drive/STAT6000/ResearchProject/train2")
pods_metadata = MetadataCatalog.get("pods_train")
dataset_dicts = DatasetCatalog.get("pods_train")
for d in random.sample(dataset_dicts, 3):
img = cv2.imread(d["file_name"])
visualizer = Visualizer(img[:, :, ::-1], metadata=pods_metadata, scale=0.5)
out = visualizer.draw_dataset_dict(d)
cv2_imshow(out.get_image()[:, :, ::-1])
from detectron2.engine import DefaultTrainer
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("pods_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.0025
cfg.SOLVER.MAX_ITER = 300
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128
j = 0
categories = {
1: "date",
2: "from",
3: "invoice_number",
4: "product_name",
5: "product_price",
6: "product_quantity",
7: "to",
8: "total_amount"
}
for i in output_pred_classes.__iter__():
classes.append(i.cpu().numpy())
for i in output_pred_boxes.__iter__():
coords = i.cpu().numpy()
x1 = int(coords[0])
y1 = int(coords[1])
x2 = int(coords[2])
y2 = int(coords[3])
crop_img = im[y1:y2, x1:x2]
key = int(classes[j])
cv2.imwrite("./cropped images/croppedIm_" + categories[key] + ".jpg",
crop_img)
j += 1
v = Visualizer(im[:, :, ::-1], metadata=test_metadata, scale=0.8)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2_imshow(out.get_image()[:, :, ::-1])
cv2.imwrite("predictions_detectron.jpg", out.get_image()[:, :, ::-1])
exsample_frames = rgb_train.load_frames(exsample_segment)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=8, shuffle=True)
for batch_num, (data, label) in enumerate(dataloader):
frame = data[0].to('cpu').detach().numpy().copy()
frame = frame.transpose(1, 2, 3, 0)
frame = np.squeeze(frame)
break
im = frame
cfg = get_cfg()
# add project-specific config (e.g., TensorMask) here if you're not running a model in detectron2's core library
cfg.merge_from_file(
model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-Detection/faster_rcnn_R_101_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
outputs = predictor(im)
# look at the outputs. See https://detectron2.readthedocs.io/tutorials/models.html#model-output-format for specification
print(outputs["instances"].pred_classes)
print(outputs["instances"].pred_boxes)
# We can use `Visualizer` to draw the predictions on the image.
v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imwrite("./output.jpg", v.get_image()[:, :, ::-1])
def draw_panoptic_seg_predictions(self,
frame,
panoptic_seg,
segments_info,
area_threshold=None,
alpha=0.5):
frame_visualizer = Visualizer(frame, self.metadata)
pred = _PanopticPrediction(panoptic_seg, segments_info)
if self._instance_mode == ColorMode.IMAGE_BW:
frame_visualizer.output.img = frame_visualizer._create_grayscale_image(
pred.non_empty_mask())
# draw mask for all semantic segments first i.e. "stuff"
for mask, sinfo in pred.semantic_masks():
category_idx = sinfo["category_id"]
try:
mask_color = [
x / 255 for x in self.metadata.stuff_colors[category_idx]
]
except AttributeError:
mask_color = None
frame_visualizer.draw_binary_mask(
mask,
color=mask_color,
text=self.metadata.stuff_classes[category_idx],
alpha=alpha,
area_threshold=area_threshold,
)
all_instances = list(pred.instance_masks())
if len(all_instances) == 0:
return frame_visualizer.output
# draw mask for all instances second
masks, sinfo = list(zip(*all_instances))
num_instances = len(masks)
masks_rles = mask_util.encode(
np.asarray(np.asarray(masks).transpose(1, 2, 0),
dtype=np.uint8,
order="F"))
assert len(masks_rles) == num_instances
category_ids = [x["category_id"] for x in sinfo]
detected = [
_DetectedInstance(category_ids[i],
bbox=None,
mask_rle=masks_rles[i],
color=None,
ttl=8) for i in range(num_instances)
]
colors = self._assign_colors(detected)
labels = [self.metadata.thing_classes[k] for k in category_ids]
frame_visualizer.overlay_instances(
boxes=None,
masks=masks,
labels=labels,
keypoints=None,
assigned_colors=colors,
alpha=alpha,
)
return frame_visualizer.output
def draw_instance_predictions(self, frame, predictions):
"""
Draw instance-level prediction results on an image.
Args:
frame (ndarray): an RGB image of shape (H, W, C), in the range [0, 255].
predictions (Instances): the output of an instance detection/segmentation
model. Following fields will be used to draw:
"pred_boxes", "pred_classes", "scores", "pred_masks" (or "pred_masks_rle").
Returns:
output (VisImage): image object with visualizations.
"""
frame_visualizer = Visualizer(frame, self.metadata)
num_instances = len(predictions)
if num_instances == 0:
return frame_visualizer.output
boxes = predictions.pred_boxes.tensor.numpy() if predictions.has(
"pred_boxes") else None
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes.numpy() if predictions.has(
"pred_classes") else None
keypoints = predictions.pred_keypoints if predictions.has(
"pred_keypoints") else None
if predictions.has("pred_masks"):
masks = predictions.pred_masks
# mask IOU is not yet enabled
# masks_rles = mask_util.encode(np.asarray(masks.permute(1, 2, 0), order="F"))
# assert len(masks_rles) == num_instances
else:
masks = None
detected = [
_DetectedInstance(classes[i],
boxes[i],
mask_rle=None,
color=None,
ttl=8) for i in range(num_instances)
]
colors = self._assign_colors(detected)
labels = _create_text_labels(classes, scores,
self.metadata.get("thing_classes", None))
if self._instance_mode == ColorMode.IMAGE_BW:
# any() returns uint8 tensor
frame_visualizer.output.img = frame_visualizer._create_grayscale_image(
(masks.any(dim=0) > 0).numpy() if masks is not None else None)
alpha = 0.3
else:
alpha = 0.5
frame_visualizer.overlay_instances(
boxes=None
if masks is not None else boxes, # boxes are a bit distracting
masks=masks,
labels=labels,
keypoints=keypoints,
assigned_colors=colors,
alpha=alpha,
)
return frame_visualizer.output
def train_detectron():
setup_logger()
"""
Register dataset
"""
from detectron2.data.datasets import register_coco_instances
for d in ["train", "test"]:
register_coco_instances(f"bladder_masked_{d}", {}, f"Data/Vezica/Tumor_Masked/{d}.json",
f"Data/Vezica/Tumor_Masked/{d}")
"""
Display 3 random images from the dataset
"""
import random
from detectron2.data import DatasetCatalog, MetadataCatalog
dataset_dicts = DatasetCatalog.get("bladder_masked_train")
bladder_masked_metadata = MetadataCatalog.get("bladder_masked_train")
# for d in random.sample(dataset_dicts, 1):
# img = cv2.imread(d["file_name"])
# v = Visualizer(img[:, :, ::-1], metadata=bladder_masked_metadata, scale=0.5)
# v = v.draw_dataset_dict(d)
# plt.figure(figsize = (14, 10))
# plt.imshow(cv2.cvtColor(v.get_image()[:, :, ::-1], cv2.COLOR_BGR2RGB))
# plt.show()
"""
Train
"""
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("bladder_masked_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 1
cfg.OUTPUT_DIR = "./Output/Bladder_Masked"
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
# cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg.SOLVER.IMS_PER_BATCH = 1
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 1500
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
# trainer.resume_or_load(resume=False)
trainer.train()
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
cfg.DATASETS.TEST = ("bladder_masked_test",)
predictor = DefaultPredictor(cfg)
from detectron2.utils.visualizer import ColorMode
from detectron2.evaluation.sem_seg_evaluation import SemSegEvaluator
dataset_dicts = DatasetCatalog.get("bladder_masked_test")
bladder_masked_metadata = MetadataCatalog.get("bladder_masked_test")
# evaluator = SemSegEvaluator("bladder_masked_test", False, 2)
# evaluator.evaluate()
k = 1
for d in dataset_dicts:
im = cv2.imread(d["file_name"])
outputs = predictor(im)
v = Visualizer(im[:, :, ::-1],
metadata=bladder_masked_metadata,
scale=0.8,
instance_mode=ColorMode.IMAGE_BW # remove the colors of unsegmented pixels
)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
plt.figure(figsize=(14, 10))
plt.imsave("test" + str(k) + ".png", cv2.cvtColor(v.get_image()[:, :, ::-1], cv2.COLOR_BGR2RGB))
k+=1
scale = 2.0 if args.show else 1.0
if args.source == "dataloader":
train_data_loader = build_detection_train_loader(cfg)
for batch in train_data_loader:
for per_image in batch:
# Pytorch tensor is in (C, H, W) format
img = per_image["image"].permute(1, 2, 0)
if cfg.INPUT.FORMAT == "BGR":
img = img[:, :, [2, 1, 0]]
else:
img = np.asarray(
Image.fromarray(img,
mode=cfg.INPUT.FORMAT).convert("RGB"))
visualizer = Visualizer(img, metadata=metadata, scale=scale)
target_fields = per_image["instances"].get_fields()
labels = [
metadata.thing_classes[i]
for i in target_fields["gt_classes"]
]
vis = visualizer.overlay_instances(
labels=labels,
boxes=target_fields.get("gt_boxes", None),
masks=target_fields.get("gt_masks", None),
keypoints=target_fields.get("gt_keypoints", None),
)
output(vis, str(per_image["image_id"]) + ".jpg")
else:
dicts = list(
chain.from_iterable(
"""
Test the COCO json dataset loader.
Usage:
python -m detectron2.data.datasets.coco \
path/to/json path/to/image_root dataset_name
"dataset_name" can be "coco_2014_minival_100", or other
pre-registered ones
"""
from detectron2.utils.logger import setup_logger
from detectron2.utils.visualizer import Visualizer
import detectron2.data.datasets # noqa # add pre-defined metadata
import sys
logger = setup_logger(name=__name__)
assert sys.argv[3] in DatasetCatalog.list()
meta = MetadataCatalog.get(sys.argv[3])
dicts = load_coco_json(sys.argv[1], sys.argv[2], sys.argv[3])
logger.info("Done loading {} samples.".format(len(dicts)))
dirname = "coco-data-vis"
os.makedirs(dirname, exist_ok=True)
for d in dicts:
img = np.array(Image.open(d["file_name"]))
visualizer = Visualizer(img, metadata=meta)
vis = visualizer.draw_dataset_dict(d)
fpath = os.path.join(dirname, os.path.basename(d["file_name"]))
vis.save(fpath)
def draw_instance_predictions(self, frame, predictions):
"""
Draw instance-level prediction results on an image.
Args:
frame (ndarray): an RGB image of shape (H, W, C), in the range [0, 255].
predictions (Instances): the output of an instance detection/segmentation
model. Following fields will be used to draw:
"pred_boxes", "pred_classes", "scores", "pred_masks" (or "pred_masks_rle").
Returns:
output (VisImage): image object with visualizations.
"""
frame_visualizer = Visualizer(frame, self.metadata)
num_instances = len(predictions)
if num_instances == 0:
return frame_visualizer.output
boxes = predictions.pred_boxes.tensor.numpy() if predictions.has(
"pred_boxes") else None
scores = predictions.scores if predictions.has("scores") else None
classes = predictions.pred_classes.numpy() if predictions.has(
"pred_classes") else None
keypoints = predictions.pred_keypoints if predictions.has(
"pred_keypoints") else None
colors = predictions.COLOR if predictions.has(
"COLOR") else [None] * len(predictions)
durations = predictions.ID_duration if predictions.has(
"ID_duration") else None
duration_threshold = self.metadata.get("duration_threshold", 0)
visibilities = None if durations is None else [
x > duration_threshold for x in durations
]
if predictions.has("pred_masks"):
masks = predictions.pred_masks
# mask IOU is not yet enabled
# masks_rles = mask_util.encode(np.asarray(masks.permute(1, 2, 0), order="F"))
# assert len(masks_rles) == num_instances
else:
masks = None
detected = [
_DetectedInstance(classes[i],
boxes[i],
mask_rle=None,
color=colors[i],
ttl=8) for i in range(num_instances)
]
if not predictions.has("COLOR"):
colors = self._assign_colors(detected)
labels = _create_text_labels(classes, scores,
self.metadata.get("thing_classes", None))
if self._instance_mode == ColorMode.IMAGE_BW:
# any() returns uint8 tensor
frame_visualizer.output.reset_image(
frame_visualizer._create_grayscale_image((masks.any(
dim=0) > 0).numpy() if masks is not None else None))
alpha = 0.3
else:
alpha = 0.5
labels = (
None if labels is None else
[y[0] for y in filter(lambda x: x[1], zip(labels, visibilities))]
) # noqa
assigned_colors = (
None if colors is None else
[y[0] for y in filter(lambda x: x[1], zip(colors, visibilities))]
) # noqa
frame_visualizer.overlay_instances(
boxes=None if masks is not None else
boxes[visibilities], # boxes are a bit distracting
masks=None if masks is None else masks[visibilities],
labels=labels,
keypoints=None if keypoints is None else keypoints[visibilities],
assigned_colors=assigned_colors,
alpha=alpha,
)
return frame_visualizer.output
DatasetCatalog.register(
"products_" + d,
lambda d=d: get_product_page_dicts(
"/home/kevin/bin/scraping_engine/data/export-2020-10-22T19_24_01.639Z.json",
"/home/kevin/bin/scraping_engine/data/product_images_2"))
MetadataCatalog.get("products_" + d).set(
thing_classes=["title", "rating", "ratings", "price", "description"])
products_metadata = MetadataCatalog.get("products_train")
dicts = get_product_page_dicts(
"/home/kevin/bin/scraping_engine/data/export-2020-10-22T19_24_01.639Z.json",
"/home/kevin/bin/scraping_engine/data/product_images_2")
for i, d in enumerate(random.sample(dicts, 3)):
img = cv2.imread(d["file_name"])
visualizer = Visualizer(img[:, :, ::-1],
metadata=products_metadata,
scale=0.5)
out = visualizer.draw_dataset_dict(d)
print(out)
cv2.imwrite(f"test-img{i}.png", out.get_image()[:, :, ::-1])
# DatasetCatalog.register("balloon_" + d, lambda d=d: get_balloon_dicts("balloon/" + d))
# MetadataCatalog.get("balloon_" + d).set(thing_classes=["balloon"])
balloon_metadata = MetadataCatalog.get("balloon_train")
from detectron2.engine import DefaultTrainer
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml"))
img = cv2.imread("./jua.jpg")
# img = cv2.resize(img,(0,0),fx=0.5,fy=0.5)
print(img.shape)
# %% key point
# create config for Keypoints detection
_cfg = get_cfg()
_cfg.merge_from_file(
model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
_cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
_cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")
_cfg.MODEL.DEVICE = 'cpu' # device를 cpu로 선택
# Keypoints predictor
_predictor = DefaultPredictor(_cfg)
#%%
start_tick = time.time()
outputs = _predictor(img)
print(f'delay { time.time() - start_tick }')
#%%
# using `Visualizer` to draw the predictions on the image.
v = Visualizer(cv2.cvtColor(img, cv2.COLOR_BGR2RGB),
MetadataCatalog.get(_cfg.DATASETS.TRAIN[0]),
scale=1.2)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
display(Image.fromarray(out.get_image()))
# %%
# Register data inside detectron
# With DATASET_SIZES one can limit the size of these datasets
d = "inference"
inference_dict = get_lofar_dicts(
os.path.join(DATASET_PATH, "VIA_json_inference.pkl"))
DatasetCatalog.register(d, lambda d=d: inference_dict)
MetadataCatalog.get(d).set(thing_classes=["radio_source"])
lofar_metadata = MetadataCatalog.get(d)
print("Sample and plot input data as sanity check")
#"""
for i, dic in enumerate(random.sample(inference_dict, 3)):
print(dic["file_name"])
img = imread(dic["file_name"])
visualizer = Visualizer(img[:, :, ::-1], metadata=lofar_metadata, scale=1)
vis = visualizer.draw_dataset_dict(dic)
a = vis.get_image()[:, :, ::-1]
plt.figure(figsize=(10, 10))
plt.imshow(a)
plt.savefig(
os.path.join(cfg.OUTPUT_DIR,
f"random_input_example_for_sanity_check_{i}.jpg"))
plt.close()
#"""
# Inference mode
# To implement the LOFAR relevant metrics I changed
# DefaultTrainer into LOFARTrainer
# where the latter calls LOFAREvaluator within build_hooks instead of the default evaluator
