def __init__(
self,
dataset_name,
tasks=None,
distributed=True,
output_dir=None,
*,
use_fast_impl=True,
kpt_oks_sigmas=(),
):
"""
Args:
dataset_name (str): name of the dataset to be evaluated.
It must have either the following corresponding metadata:
"json_file": the path to the COCO format annotation
Or it must be in detectron2's standard dataset format
so it can be converted to COCO format automatically.
tasks (tuple[str]): tasks that can be evaluated under the given
configuration. A task is one of "bbox", "segm", "keypoints".
By default, will infer this automatically from predictions.
distributed (True): if True, will collect results from all ranks and run evaluation
in the main process.
Otherwise, will only evaluate the results in the current process.
output_dir (str): optional, an output directory to dump all
results predicted on the dataset. The dump contains two files:
1. "instances_predictions.pth" a file that can be loaded with `torch.load` and
contains all the results in the format they are produced by the model.
2. "coco_instances_results.json" a json file in COCO's result format.
use_fast_impl (bool): use a fast but **unofficial** implementation to compute AP.
Although the results should be very close to the official implementation in COCO
API, it is still recommended to compute results with the official API for use in
papers. The faster implementation also uses more RAM.
kpt_oks_sigmas (list[float]): The sigmas used to calculate keypoint OKS.
See http://cocodataset.org/#keypoints-eval
When empty, it will use the defaults in COCO.
Otherwise it should be the same length as ROI_KEYPOINT_HEAD.NUM_KEYPOINTS.
"""
self._logger = logging.getLogger(__name__)
self._distributed = distributed
self._output_dir = output_dir
self._use_fast_impl = use_fast_impl
if tasks is not None and isinstance(tasks, CfgNode):
kpt_oks_sigmas = (
tasks.TEST.KEYPOINT_OKS_SIGMAS if not kpt_oks_sigmas else kpt_oks_sigmas
)
self._logger.warn(
"COCO Evaluator instantiated using config, this is deprecated behavior."
" Please pass in explicit arguments instead."
)
self._tasks = None # Infering it from predictions should be better
else:
self._tasks = tasks
self._cpu_device = torch.device("cpu")
self._metadata = MetadataCatalog.get(dataset_name)
if not hasattr(self._metadata, "json_file"):
self._logger.info(
f"'{dataset_name}' is not registered by `register_coco_instances`."
" Therefore trying to convert it to COCO format ..."
)
cache_path = os.path.join(output_dir, f"{dataset_name}_coco_format.json")
self._metadata.json_file = cache_path
convert_to_coco_json(dataset_name, cache_path)
json_file = PathManager.get_local_path(self._metadata.json_file)
with contextlib.redirect_stdout(io.StringIO()):
self._coco_api = COCO(json_file)
# Test set json files do not contain annotations (evaluation must be
# performed using the COCO evaluation server).
self._do_evaluation = "annotations" in self._coco_api.dataset
if self._do_evaluation:
self._kpt_oks_sigmas = kpt_oks_sigmas
train_folder = '../../Others/flownet2/flownet2-pytorch/KAIST_img/KAIST_train_flow/'
#train_json_path = '../../../Datasets/'+dataset+'/train/thermal_annotations_4class.json'
train_json_path = '../../../Datasets/'+dataset+'/train/KAIST_train_flow_annotation.json'
#train_json_path = '../../../Datasets/'+dataset+'/train/thermal_annotations.json'
# Validation path
#val_folder = '../../../Datasets/FLIR/test/
val_folder = '../../Others/flownet2/flownet2-pytorch/KAIST_img/KAIST_test_flow/'
#val_json_path = '../../../Datasets/'+dataset+'/val/thermal_annotations_4class.json'
val_json_path = '../../../Datasets/'+dataset+'/test/KAIST_test_flow_annotation.json'
print('train json path:', train_json_path)
print('test json path:', val_json_path)
# Register dataset
dataset_train = 'KAIST_train'
register_coco_instances(dataset_train, {}, train_json_path, train_folder)
FLIR_metadata_train = MetadataCatalog.get(dataset_train)
dataset_dicts_train = DatasetCatalog.get(dataset_train)
# Test on validation set
dataset_test = 'KAIST_test'
register_coco_instances(dataset_test, {}, val_json_path, val_folder)
FLIR_metadata_test = MetadataCatalog.get(dataset_test)
dataset_dicts_test = DatasetCatalog.get(dataset_test)
model = 'faster_rcnn_R_101_FPN_3x'
#files_names = [f for f in listdir(train_path) if isfile(join(train_path, f))]
out_folder = 'out_training/KAIST_flow_0404/'
out_model_path = os.path.join(out_folder, 'out_model_final.pth')
if not os.path.exists(out_folder):
for bbox, objectness_logit in zip(bboxes, objectness_logits):
proposal_by_image[image_id].append({
'bbox':
bbox,
'score':
1 / (1 + np.exp(objectness_logit))
})
with PathManager.open(args.result, "r") as f:
predictions = json.load(f)
pred_by_image = defaultdict(list)
for p in predictions:
pred_by_image[p["image_id"]].append(p)
dicts = list(DatasetCatalog.get(args.dataset))
metadata = MetadataCatalog.get(args.dataset)
if hasattr(metadata, "thing_dataset_id_to_contiguous_id"):
def dataset_id_map(ds_id):
return metadata.thing_dataset_id_to_contiguous_id[ds_id]
elif "lvis" in args.dataset:
# LVIS results are in the same format as COCO results, but have a different
# mapping from dataset category id to contiguous category id in [0, #categories - 1]
def dataset_id_map(ds_id):
return ds_id - 1
else:
raise ValueError("Unsupported dataset: {}".format(args.dataset))
os.makedirs(args.output, exist_ok=True)
def register_pseudo(name, class_names=CLASS_NAMES):
DatasetCatalog.register(name,
lambda: load_personface_instances(class_names))
MetadataCatalog.get(name).set(thing_classes=list(class_names))
from detectron2.data.datasets import register_coco_instances
from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.utils.visualizer import Visualizer
from detectron2.config import get_cfg
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor, DefaultTrainer
from detectron2.utils.visualizer import ColorMode
import os
import cv2
import random
register_coco_instances('pcb', {}, "person_car_bus_val2017.json",
"../cocoData/val2017")
pcb_metadata = MetadataCatalog.get("pcb")
pcb_data = DatasetCatalog.get("pcb")
f = DatasetCatalog._REGISTERED['pcb']
print(f)
# for d in random.sample(pcb_data,3):
# img = cv2.imread(d["file_name"])
# visualizer = Visualizer(img[:,:,::-1],metadata=pcb_metadata,scale=0.5)
# vis = visualizer.draw_dataset_dict(d)
# cv2.imshow("image", vis.get_image()[:, :, ::-1])
# cv2.waitKey(0)
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
pycocotools.mask.encode(
b_a
), #cfg.INPUT.MASK_FORMAT must be set to bitmask if using the default data loader with such format.
"category_id":
classes.index(imgclass),
# "category_id": 0,
}]
dataset_dicts.append(record)
return dataset_dicts # Returns a dict of all images with their respective descriptions
for d in ["train", "test"]:
DatasetCatalog.register("logo_" + d,
lambda d=d: get_logos(path + d + "set.txt"))
MetadataCatalog.get("logo_" + d).set(thing_classes=classes)
logo_metadata = MetadataCatalog.get("logo_train")
dataset_dicts = DatasetCatalog.get("logo_train")
# Take three random images, apply the masks and output them
# for d in random.sample(dataset_dicts, 3):
# img = cv2.imread(d["file_name"])
# visualizer = Visualizer(img[:, :, ::-1], metadata=logo_metadata)
# out = visualizer.draw_dataset_dict(d)
# result = out.get_image()[:, :, ::-1]
# cv2.imwrite("output/output%s.jpg" % str(d["file_name"][-8:-6]), result)
model = "COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"
# model = "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"
def tearDown(self):
# Need to remove injected dataset
injected_dataset = set(DatasetCatalog) - self._builtin_datasets
for ds in injected_dataset:
DatasetCatalog.remove(ds)
MetadataCatalog.remove(ds)
def load_coco_json(json_file,
image_root,
dataset_name=None,
extra_annotation_keys=None):
"""
Load a json file with COCO's instances annotation format.
Currently supports instance detection, instance segmentation,
and person keypoints annotations.
Args:
json_file (str): full path to the json file in COCO instances annotation format.
image_root (str): the directory where the images in this json file exists.
dataset_name (str): the name of the dataset (e.g., coco_2017_train).
If provided, this function will also put "thing_classes" into
the metadata associated with this dataset.
extra_annotation_keys (list[str]): list of per-annotation keys that should also be
loaded into the dataset dict (besides "iscrowd", "bbox", "keypoints",
"category_id", "segmentation"). The values for these keys will be returned as-is.
For example, the densepose annotations are loaded in this way.
Returns:
list[dict]: a list of dicts in Detectron2 standard format. (See
`Using Custom Datasets </tutorials/datasets.html>`_ )
Notes:
1. This function does not read the image files.
The results do not have the "image" field.
"""
from pycocotools.coco import COCO
timer = Timer()
json_file = PathManager.get_local_path(json_file)
with contextlib.redirect_stdout(io.StringIO()):
coco_api = COCO(json_file)
if timer.seconds() > 1:
logger.info("Loading {} takes {:.2f} seconds.".format(
json_file, timer.seconds()))
id_map = None
if dataset_name is not None:
meta = MetadataCatalog.get(dataset_name)
cat_ids = sorted(coco_api.getCatIds())
cats = coco_api.loadCats(cat_ids)
# The categories in a custom json file may not be sorted.
thing_classes = [
c["name"] for c in sorted(cats, key=lambda x: x["id"])
]
meta.thing_classes = thing_classes
# In COCO, certain category ids are artificially removed,
# and by convention they are always ignored.
# We deal with COCO's id issue and translate
# the category ids to contiguous ids in [0, 80).
# It works by looking at the "categories" field in the json, therefore
# if users' own json also have incontiguous ids, we'll
# apply this mapping as well but print a warning.
if not (min(cat_ids) == 1 and max(cat_ids) == len(cat_ids)):
if "coco" not in dataset_name:
logger.warning("""
Category ids in annotations are not in [1, #categories]! We'll apply a mapping for you.
""")
id_map = {v: i for i, v in enumerate(cat_ids)}
meta.thing_dataset_id_to_contiguous_id = id_map
# sort indices for reproducible results
img_ids = sorted(list(coco_api.imgs.keys()))
# imgs is a list of dicts, each looks something like:
# {'license': 4,
# 'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',
# 'file_name': 'COCO_val2014_000000001268.jpg',
# 'height': 427,
# 'width': 640,
# 'date_captured': '2013-11-17 05:57:24',
# 'id': 1268}
imgs = coco_api.loadImgs(img_ids)
# anns is a list[list[dict]], where each dict is an annotation
# record for an object. The inner list enumerates the objects in an image
# and the outer list enumerates over images. Example of anns[0]:
# [{'segmentation': [[192.81,
# 247.09,
# ...
# 219.03,
# 249.06]],
# 'area': 1035.749,
# 'iscrowd': 0,
# 'image_id': 1268,
# 'bbox': [192.81, 224.8, 74.73, 33.43],
# 'category_id': 16,
# 'id': 42986},
# ...]
anns = [coco_api.imgToAnns[img_id] for img_id in img_ids]
if "minival" not in json_file:
# The popular valminusminival & minival annotations for COCO2014 contain this bug.
# However the ratio of buggy annotations there is tiny and does not affect accuracy.
# Therefore we explicitly white-list them.
ann_ids = [
ann["id"] for anns_per_image in anns for ann in anns_per_image
]
assert len(set(ann_ids)) == len(
ann_ids), "Annotation ids in '{}' are not unique!".format(
json_file)
imgs_anns = list(zip(imgs, anns))
logger.info("Loaded {} images in COCO format from {}".format(
len(imgs_anns), json_file))
dataset_dicts = []
ann_keys = ["iscrowd", "bbox", "keypoints", "category_id"
] + (extra_annotation_keys or [])
num_instances_without_valid_segmentation = 0
for (img_dict, anno_dict_list) in imgs_anns:
record = {}
record["file_name"] = os.path.join(image_root, img_dict["file_name"])
record["height"] = img_dict["height"]
record["width"] = img_dict["width"]
image_id = record["image_id"] = img_dict["id"]
objs = []
for anno in anno_dict_list:
# Check that the image_id in this annotation is the same as
# the image_id we're looking at.
# This fails only when the data parsing logic or the annotation file is buggy.
# The original COCO valminusminival2014 & minival2014 annotation files
# actually contains bugs that, together with certain ways of using COCO API,
# can trigger this assertion.
assert anno["image_id"] == image_id
assert anno.get("ignore", 0) == 0
obj = {key: anno[key] for key in ann_keys if key in anno}
segm = anno.get("segmentation", None)
if segm: # either list[list[float]] or dict(RLE)
if not isinstance(segm, dict):
# filter out invalid polygons (< 3 points)
segm = [
poly for poly in segm
if len(poly) % 2 == 0 and len(poly) >= 6
]
if len(segm) == 0:
num_instances_without_valid_segmentation += 1
continue # ignore this instance
obj["segmentation"] = segm
keypts = anno.get("keypoints", None)
if keypts: # list[int]
for idx, v in enumerate(keypts):
if idx % 3 != 2:
# COCO's segmentation coordinates are floating points in [0, H or W],
# but keypoint coordinates are integers in [0, H-1 or W-1]
# Therefore we assume the coordinates are "pixel indices" and
# add 0.5 to convert to floating point coordinates.
keypts[idx] = v + 0.5
obj["keypoints"] = keypts
obj["bbox_mode"] = BoxMode.XYWH_ABS
if id_map:
obj["category_id"] = id_map[obj["category_id"]]
objs.append(obj)
record["annotations"] = objs
dataset_dicts.append(record)
if num_instances_without_valid_segmentation > 0:
logger.warn(
"Filtered out {} instances without valid segmentation. "
"There might be issues in your dataset generation process.".format(
num_instances_without_valid_segmentation))
return dataset_dicts
def convert_to_coco_dict(dataset_name):
"""
Convert a dataset in detectron2's standard format into COCO json format
Generic dataset description can be found here:
https://detectron2.readthedocs.io/tutorials/datasets.html#register-a-dataset
COCO data format description can be found here:
http://cocodataset.org/#format-data
Args:
dataset_name:
name of the source dataset
must be registered in DatastCatalog and in detectron2's standard format
Returns:
coco_dict: serializable dict in COCO json format
"""
dataset_dicts = DatasetCatalog.get(dataset_name)
categories = [{
"id": id,
"name": name
} for id, name in enumerate(
MetadataCatalog.get(dataset_name).thing_classes)]
logger.info("Converting dataset dicts into COCO format")
coco_images = []
coco_annotations = []
for image_id, image_dict in enumerate(dataset_dicts):
coco_image = {
"id": image_dict.get("image_id", image_id),
"width": image_dict["width"],
"height": image_dict["height"],
"file_name": image_dict["file_name"],
}
coco_images.append(coco_image)
anns_per_image = image_dict["annotations"]
for annotation in anns_per_image:
# create a new dict with only COCO fields
coco_annotation = {}
# COCO requirement: XYWH box format
bbox = annotation["bbox"]
bbox_mode = annotation["bbox_mode"]
bbox = BoxMode.convert(bbox, bbox_mode, BoxMode.XYWH_ABS)
# COCO requirement: instance area
if "segmentation" in annotation:
# Computing areas for instances by counting the pixels
segmentation = annotation["segmentation"]
# TODO: check segmentation type: RLE, BinaryMask or Polygon
polygons = PolygonMasks([segmentation])
area = polygons.area()[0].item()
else:
# Computing areas using bounding boxes
bbox_xy = BoxMode.convert(bbox, BoxMode.XYWH_ABS,
BoxMode.XYXY_ABS)
area = Boxes([bbox_xy]).area()[0].item()
if "keypoints" in annotation:
keypoints = annotation["keypoints"] # list[int]
for idx, v in enumerate(keypoints):
if idx % 3 != 2:
# COCO's segmentation coordinates are floating points in [0, H or W],
# but keypoint coordinates are integers in [0, H-1 or W-1]
# For COCO format consistency we substract 0.5
# https://github.com/facebookresearch/detectron2/pull/175#issuecomment-551202163
keypoints[idx] = v - 0.5
if "num_keypoints" in annotation:
num_keypoints = annotation["num_keypoints"]
else:
num_keypoints = sum(kp > 0 for kp in keypoints[2::3])
# COCO requirement:
# linking annotations to images
# "id" field must start with 1
coco_annotation["id"] = len(coco_annotations) + 1
coco_annotation["image_id"] = coco_image["id"]
coco_annotation["bbox"] = [round(float(x), 3) for x in bbox]
coco_annotation["area"] = area
coco_annotation["category_id"] = annotation["category_id"]
coco_annotation["iscrowd"] = annotation.get("iscrowd", 0)
# Add optional fields
if "keypoints" in annotation:
coco_annotation["keypoints"] = keypoints
coco_annotation["num_keypoints"] = num_keypoints
if "segmentation" in annotation:
coco_annotation["segmentation"] = annotation["segmentation"]
coco_annotations.append(coco_annotation)
logger.info(
"Conversion finished, "
f"num images: {len(coco_images)}, num annotations: {len(coco_annotations)}"
)
info = {
"date_created": str(datetime.datetime.now()),
"description":
"Automatically generated COCO json file for Detectron2.",
}
coco_dict = {
"info": info,
"images": coco_images,
"annotations": coco_annotations,
"categories": categories,
"licenses": None,
}
return coco_dict
def draw_output(self, outputs, im):
# Get class labels found in image
metadata = MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0])
# Draw segmentations on image
v = Visualizer(im[:, :, ::-1], metadata, scale=1)
return v.draw_instance_predictions(outputs["instances"].to("cpu"))
cfg = get_cfg()
cfg.MODEL.DEVICE='cpu'
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.7 # set threshold for this model
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml")
cap = cv2.VideoCapture(0)
while(True):
# Capture frame-by-frame
ret, frame = cap.read()
# Our operations on the frame come here
predictor = DefaultPredictor(cfg)
outputs = predictor(frame)
v = Visualizer(frame[:,:,::-1], MetadataCatalog.get(cfg.DATASETS.TRAIN[0]), scale=1.2)
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
# Display the resulting frame
cv2.imshow('frame',v.get_image()[:, :, ::-1])
# Write the output frame to disk
#writer.write(v.get_image()[:, :, ::-1])
if cv2.waitKey(1) & 0xFF == ord('x'):
break
# When everything done, release the capture
cap.release()
cv2.destroyAllWindows()
ctypes.windll.user32.MessageBoxW(0, "The video has been succefully completed", "Finished", 1)
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog
from detectron2.data.datasets import register_coco_instances
from detectron2.utils.visualizer import ColorMode
from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
# Register the dataset
register_coco_instances("train_set", {}, "/home/labuser/ros_ws/src/odhe_ros/scene_camera_dataset/train/annotations.json", "/home/labuser/ros_ws/src/odhe_ros/scene_camera_dataset/train")
register_coco_instances("test_set", {}, "/home/labuser/ros_ws/src/odhe_ros/scene_camera_dataset/test/annotations.json", "/home/labuser/ros_ws/src/odhe_ros/scene_camera_dataset/test")
# train_metadata = MetadataCatalog.get("train_set")
# print(train_metadata)
# dataset_dicts = DatasetCatalog.get("train_set")
test_metadata = MetadataCatalog.get("test_set")
print(test_metadata)
dataset_dicts_test = DatasetCatalog.get("test_set")
# # Visualize Training data
# for d in random.sample(dataset_dicts, 3):
# img = cv2.imread(d["file_name"])
# visualizer = Visualizer(img[:, :, ::-1], metadata=train_metadata, scale=0.5)
# vis = visualizer.draw_dataset_dict(d)
# cv2.imshow('Training Image', vis.get_image()[:, :, ::-1])
# cv2.waitKey(0)
# # Visualize Test data
# for d in random.sample(dataset_dicts_test, 3):
# img = cv2.imread(d["file_name"])
# visualizer = Visualizer(img[:, :, ::-1], metadata=test_metadata, scale=0.5)
dets.append(det)
# ann = ann[['ImageID', 'LabelName', 'XMin', 'XMax', 'YMin', 'YMax']].values
# det = det[['ImageID', 'LabelName', 'Conf', 'XMin', 'XMax', 'YMin', 'YMax']].values
mean_ap, average_precisions = mean_average_precision_for_boxes(anns, dets)
if __name__ == "__main__":
args = parser.parse_args()
split = 'test_6'
dataset_name = "mtsd_" + split
DatasetCatalog.register(dataset_name,
lambda d=split: utils.load_obj(split))
MetadataCatalog.get(dataset_name).set(thing_classes=CATEGORIES)
Meta_data = MetadataCatalog.get(dataset_name)
cfg = get_cfg()
cfg.merge_from_file(
model_zoo.get_config_file(
"COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml"))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.0 # set threshold for this model
cfg.MODEL.WEIGHTS = args.path_to_model
cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(CATEGORIES)
cfg.merge_from_file('configs/small_objects.yaml')
predictor = DefaultPredictor(cfg)
validate(cfg, Meta_data, split)
"category_id": int(v[2]) - 1,
"iscrowd": 0
}
objs.append(deepcopy(obj))
prev_img_id = 10000 * number_of_folder + number_of_frame
print("Done: ", ann)
return dataset_dicts
from detectron2.data import DatasetCatalog, MetadataCatalog
#dataset = get_KITTI_dicts("datasets/KITTI")
for d in ["train"]:
DatasetCatalog.register("KITTI_" + d, lambda d=d: get_KITTI_dicts("datasets/KITTI/"))
MetadataCatalog.get("KITTI_" + d).set(thing_classes=["car", "pedestrian"])
KITTI_metadata = MetadataCatalog.get("KITTI_train")
class Trainer(DefaultTrainer):
"""
This is the same Trainer except that we rewrite the
`build_train_loader` method.
"""
def __init__(self, cfg):
"""
Args:
cfg (CfgNode):
Use the custom checkpointer, which loads other backbone models
with matching heuristics.
from detectron2.evaluation import COCOEvaluator, inference_on_dataset, LVISEvaluator
from detectron2.data import build_detection_test_loader
from detectron2.engine import DefaultTrainer
import matplotlib
from detectron2.utils.visualizer import ColorMode
from detectron2.structures import BoxMode
with open(r'E:\PyProjects\Faster-RCNN\Database\data24\val\data_test.txt',
'r') as in_file:
testset_dicts = json.load(in_file)
for dict in testset_dicts:
for ann in dict["annotations"]:
ann["bbox_mode"] = BoxMode.XYXY_ABS
DatasetCatalog.register("mol_val", lambda d="val": testset_dicts)
MetadataCatalog.get("mol_val").set(thing_classes=["L", "R"])
mol_metadata = MetadataCatalog.get("mol_val")
cfg = get_cfg()
# cfg.MODEL.DEVICE = "cpu"
cfg.merge_from_file(
model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TEST = ("mol_val", )
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 2
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.3 # set threshold for this model
cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST = 0.3
cfg.MODEL.RPN.NMS_THRESH = 0.8
cfg.MODEL.RPN.IOU_THRESHOLDS = [0.6, 0.8]
# Find a model from detectron2's model zoo. You can use the https://dl.fbaipublicfiles... url as well
cfg.MODEL.WEIGHTS = r"E:\PyProjects\Faster-RCNN\Results\data24\output\model_final.pth"
cfg.MODEL.ANCHOR_GENERATOR.SIZES = [[8, 16, 32, 64, 128]]
COCOEvaluator,
COCOPanopticEvaluator,
DatasetEvaluators,
LVISEvaluator,
PascalVOCDetectionEvaluator,
SemSegEvaluator,
verify_results,
)
from detectron2.modeling import GeneralizedRCNNWithTTA
from detectron2.data.datasets import register_coco_instances
from detectron2.data import MetadataCatalog, DatasetCatalog
register_coco_instances("plans_train", {}, "./Flo2PlanAll-Seg/annotations/flo2plan_instances_final.json", "./Flo2PlanAll-Seg/images")
register_coco_instances("plans_test", {}, "./Flo2PlanAll-Seg/annotations/flo2plan_instances_final.json", "./Flo2PlanAll-Seg/images")
plans_metadata = MetadataCatalog.get("plans_train")
dataset_dicts = DatasetCatalog.get("plans_train")
#dataset_dicts_test = DatasetCatalog.get("plans_test")
class Trainer(DefaultTrainer):
"""
We use the "DefaultTrainer" which contains a number pre-defined logic for
standard training workflow. They may not work for you, especially if you
are working on a new research project. In that case you can use the cleaner
"SimpleTrainer", or write your own training loop.
"""
@classmethod
def build_evaluator(cls, cfg, dataset_name, output_folder=None):
"/home/group00/mcv/datasets/MOTSChallenge/train/instances_txt/0005.txt")
train_paths.append(
"/home/group00/mcv/datasets/MOTSChallenge/train/instances_txt/0009.txt")
train_paths.append(
"/home/group00/mcv/datasets/MOTSChallenge/train/instances_txt/0011.txt")
OUTPUT_DIR = "/home/group00/working/week4/model_evaluation"
if not os.path.exists(OUTPUT_DIR):
os.makedirs(OUTPUT_DIR)
#train dataset
for d in [train_paths]:
DatasetCatalog.register('train_kitti-mots',
lambda d=d: get_kitti_dataset_train(d))
MetadataCatalog.get('train_kitti-mots').set(
thing_classes=['Person', 'NA', 'Car'])
#val dataset
for d in [val_paths]:
DatasetCatalog.register('val_kitti-mots',
lambda d=d: get_kitti_dataset_val(d))
MetadataCatalog.get('val_kitti-mots').set(
thing_classes=['Person', 'NA', 'Car'])
# Pre-trained
print('Loading pre-trained models...')
cfg = get_cfg()
#Select model
cfg.merge_from_file(
if __name__ == "__main__":
args = parse_args()
model_path = os.path.join(args.config, args.model + '.yaml')
print('[INFO] Using model: ', model_path)
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(model_path))
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5 # set threshold for this model
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_path)
cfg.OUTPUT_DIR = os.path.join(args.output, args.config, args.model)
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
metadata = MetadataCatalog.get(cfg.DATASETS.TRAIN[0])
thing_classes = metadata.thing_classes
map_classes = {1: 2, 2: 0}
dataset = 'KITTI-MOTS'
for d in ['train', 'test']:
DatasetCatalog.register(
dataset + '_' + d,
lambda d=d: load_dataset(d, thing_classes, map_classes))
MetadataCatalog.get(dataset + '_' + d).set(thing_classes=thing_classes)
metadata = MetadataCatalog.get(dataset + '_train')
cfg.DATASETS.TRAIN = (dataset + '_train', )
cfg.DATASETS.TEST = (dataset + '_test', )
cfg.DATALOADER.NUM_WORKERS = 2
"category_id": 0,
"iscrowd": 0
}
objs.append(obj)
record["annotations"] = objs
dataset_dicts.append(record)
return dataset_dicts
if __name__ == '__main__':
for d in ["train", "val"]:
DatasetCatalog.register(
"./images/balloon/" + d,
lambda d=d: get_balloon_dicts("./images/balloon/" + d))
MetadataCatalog.get("./images/balloon/" +
d).set(thing_classes=["balloon"])
balloon_metadata = MetadataCatalog.get("./images/balloon/train")
cfg = get_cfg()
cfg.merge_from_file(
"./configs/mask_rcnn_R_50_FPN_3x.yaml"
) #./detectron2_repo/configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml
cfg.DATASETS.TRAIN = ("./images/balloon/train", )
cfg.DATASETS.TEST = () # no metrics implemented for this dataset
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = "detectron2://COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x/137849600/model_final_f10217.pkl" # initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 300 # 300 iterations seems good enough, but you can certainly train longer
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128 # faster, and good enough for this toy dataset
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1 # only has one class (ballon)
import annotations
from config import Config
from detectron2.data import DatasetCatalog, MetadataCatalog
oyster_cfg = Config(0)
for d in ["train", "val"]:
dataset = annotations.labelmeDict(oyster_cfg.folders['data'], d)
# annotations.draw(dataset)
DatasetCatalog.register(
"oyster_" + d,
lambda d=d: annotations.labelmeDict(oyster_cfg.folders['data'], d))
MetadataCatalog.get("oyster_" + d).set(thing_classes=["oyster"])
oyster_metadata = MetadataCatalog.get("oyster_" + d)
annotations.draw_masks(dataset, oyster_metadata)
obj = {
"bbox": [np.min(px), np.min(py), np.max(px), np.max(py)],
"bbox_mode": BoxMode.XYXY_ABS,
"segmentation": [poly],
"category_id": 0,
"iscrowd": 0
}
objs.append(obj)
record["annotations"] = objs
dataset_dicts.append(record)
return dataset_dicts
from detectron2.data import DatasetCatalog, MetadataCatalog
for d in ["train", "val"]:
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")
print('finish')
dataset_dicts = get_balloon_dicts("balloon/train")
for d in random.sample(dataset_dicts,1 ):
img = cv2.imread(d["file_name"])
visualizer = Visualizer(img[:, :, ::-1], metadata=balloon_metadata, scale=0.125)
vis = visualizer.draw_dataset_dict(d)
plt.imshow(vis.get_image()[:, :, ::-1])
print('finish')
from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg
model_file_path = model_zoo.get_config_file(
"COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml")
model_weights = model_zoo.get_checkpoint_url(
"COCO-Detection/faster_rcnn_X_101_32x8d_FPN_3x.yaml")
# 创建一个detectron2配置
cfg = get_cfg()
cfg.MODEL.DEVICE = 'cpu'
# 要创建的模型的名称
cfg.merge_from_file(model_file_path)
# 为模型设置阈值
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
# 加载模型需要的参数数据
cfg.MODEL.WEIGHTS = model_weights
# 基于配置创建一个预测器
predictor = DefaultPredictor(cfg)
# 利用这个预测器对准备好的图像进行分析并得到结果
outputs = predictor(img)
# 使用Visualizer对结果可视化，img[:, :, ::-1]实现RGB到BGR通道的相互转换，scale表示输出图像的缩放尺度
v = Visualizer(img[:, :, ::-1],
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
scale=0.8)
# 将检测到的物体标注在图像上
v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
# 获得绘制的图像
result = v.get_image()[:, :, ::-1]
# 将影像保存到文件
cv2.imshow('result', result)
cv2.waitKey(0)
def build_categories(configs):
MetadataCatalog.get('inference').set(
thing_classes=configs['Detectron2']['LABEL_LIST']['kfashion'])
fashion_metadata = MetadataCatalog.get('inference')
return fashion_metadata
os.path.join('dataset', 'annotations', 'instances_train2014.json'))
val_annotations = COCO(
os.path.join('dataset', 'annotations', 'instances_val2014.json'))
clsMap = getId2ClassMap()
gt = json.load(open('dataset/annotations/instances_train2014.json', 'r'))
cocoDict = {}
cocoDict['info'] = gt['info']
cocoDict['licenses'] = gt['licenses']
cocoDict['images'] = []
cocoDict['annotations'] = []
cocoDict['categories'] = gt['categories']
classMapping = MetadataCatalog.get(
cfg.DATASETS.TRAIN[0]).thing_dataset_id_to_contiguous_id
rev_cls_map = {v: k for k, v in classMapping.items()}
dataset = FasterRCNNDataset(data_folder, data_name, 'TEST', transform=None)
preds = {}
annId = 0
for i, (img_id, img, img_h, img_w) in tqdm(enumerate(dataset),
total=len(dataset)):
if i % 5 != 0:
continue
outputs = predictor(255 * img.permute(1, 2, 0).numpy())
imgInfo = {
"license": 4,
"file_name": "",
def detect(video_path):
save_visual_detections = False
results_dir = 'results/task1_1/retina'
coco_car_id = 2
model = 'retinanet_R_50_FPN_3x'
model_path = 'COCO-Detection/' + model + '.yaml'
print(model_path)
# Run a pre-trained detectron2 model
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file(model_path))
cfg.MODEL.RETINANET.SCORE_THRESH_TEST = 0.5
cfg.MODEL.RETINANET.NMS_THRESH_TEST = 0.4
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(model_path)
cfg.OUTPUT_DIR = results_dir
os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
predictor = DefaultPredictor(cfg)
det_path = os.path.join(cfg.OUTPUT_DIR, 'detections.txt')
if os.path.exists(det_path):
os.remove(det_path)
vidcap = cv2.VideoCapture(video_path)
num_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
# num_frames = 3
times = []
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
for frame_id in tqdm(range(num_frames)):
_, frame = vidcap.read()
start.record()
outputs = predictor(frame)
end.record()
torch.cuda.synchronize()
times.append(start.elapsed_time(end))
pred_boxes = outputs["instances"].pred_boxes.to("cpu")
scores = outputs["instances"].scores.to("cpu")
pred_classes = outputs["instances"].pred_classes.to("cpu")
for idx, pred in enumerate(pred_classes):
if pred.item() == coco_car_id:
box = pred_boxes[idx].tensor.numpy()[0]
# Format: <frame>, <id>, <bb_left>, <bb_top>, <bb_width>, <bb_height>, <conf>, <x>, <y>, <z>
det = str(frame_id + 1) + ',-1,' + str(box[0]) + ',' + str(
box[1]) + ',' + str(box[2] - box[0]) + ',' + str(
box[3] - box[1]) + ',' + str(
scores[idx].item()) + ',-1,-1,-1\n'
with open(det_path, 'a') as f:
f.write(det)
if save_visual_detections:
output_path = os.path.join(cfg.OUTPUT_DIR,
'det_frame_' + str(frame_id) + '.png')
v = Visualizer(frame[:, :, ::-1],
MetadataCatalog.get(cfg.DATASETS.TRAIN[0]),
scale=1)
out = v.draw_instance_predictions(outputs["instances"].to("cpu"))
cv2.imwrite(output_path, out.get_image()[:, :, ::-1])
print('Inference time (s/img): ', np.mean(times) / 1000)
return det_path
def display_sample(self,
rgb_obs,
semantic_obs,
depth_obs,
mainobj=None,
visualize=False):
rgb_img = Image.fromarray(rgb_obs, mode="RGBA")
semantic_img = Image.new(
"P", (semantic_obs.shape[1], semantic_obs.shape[0]))
semantic_img.putpalette(d3_40_colors_rgb.flatten())
semantic_img.putdata((semantic_obs.flatten() % 40).astype(np.uint8))
semantic_img = semantic_img.convert("RGBA")
# st()
depth_img = Image.fromarray((depth_obs / 10 * 255).astype(np.uint8),
mode="L")
display_img = cv2.cvtColor(np.asarray(rgb_img), cv2.COLOR_RGB2BGR)
#print(display_img.shape)
# mask_image = False
# if mask_image and mainobj is not None:
# main_id = int(mainobj.id[1:])
# print("MAINID ", main_id)
# # semantic = observations["semantic_sensor"]
# display_img[semantic_obs == main_id] = [1, 0, 1]
# st()
#display_img = cv2
plt.imshow(display_img)
plt.show()
im = rgb_img[..., :3]
im = im[:, :, ::-1]
outputs = self.maskrcnn(im)
pred_masks = outputs['instances'].pred_masks
pred_boxes = outputs['instances'].pred_boxes.tensor
pred_classes = outputs['instances'].pred_classes
pred_scores = outputs['instances'].scores
# converts instance segmentation to individual masks and bbox
# visualisations
v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(self.cfg.DATASETS.TRAIN[0]),
scale=1.2)
out = v.draw_instance_predictions(outputs['instances'].to("cpu"))
seg_im = out.get_image()
# cv2.imshow('img',display_img)
if visualize:
arr = [rgb_img, semantic_img, depth_img, seg_im]
titles = ['rgb', 'semantic', 'depth', 'seg_im']
plt.figure(figsize=(12, 8))
for i, data in enumerate(arr):
ax = plt.subplot(1, 3, i + 1)
ax.axis('off')
ax.set_title(titles[i])
plt.imshow(data)
# plt.pause()
plt.show()
def register_coco_panoptic_separated(name, metadata, image_root, panoptic_root,
panoptic_json, sem_seg_root,
instances_json):
"""
Register a COCO panoptic segmentation dataset named `name`.
The annotations in this registered dataset will contain both instance annotations and
semantic annotations, each with its own contiguous ids. Hence it's called "separated".
It follows the setting used by the PanopticFPN paper:
1. The instance annotations directly come from polygons in the COCO
instances annotation task, rather than from the masks in the COCO panoptic annotations.
The two format have small differences:
Polygons in the instance annotations may have overlaps.
The mask annotations are produced by labeling the overlapped polygons
with depth ordering.
2. The semantic annotations are converted from panoptic annotations, where
all "things" are assigned a semantic id of 0.
All semantic categories will therefore have ids in contiguous
range [1, #stuff_categories].
This function will also register a pure semantic segmentation dataset
named ``name + '_stuffonly'``.
Args:
name (str): the name that identifies a dataset,
e.g. "coco_2017_train_panoptic"
metadata (dict): extra metadata associated with this dataset.
image_root (str): directory which contains all the images
panoptic_root (str): directory which contains panoptic annotation images
panoptic_json (str): path to the json panoptic annotation file
sem_seg_root (str): directory which contains all the ground truth segmentation annotations.
instances_json (str): path to the json instance annotation file
"""
panoptic_name = name + "_separated"
logger.info("[MoDet] Loaded COCO Panoptic: {}, Instance: {}".format(
panoptic_json, instances_json))
DatasetCatalog.register(
panoptic_name,
lambda: merge_to_panoptic(
load_coco_json(instances_json, image_root, panoptic_name),
load_sem_seg(sem_seg_root, image_root),
),
)
MetadataCatalog.get(panoptic_name).set(
panoptic_root=panoptic_root,
image_root=image_root,
panoptic_json=panoptic_json,
sem_seg_root=sem_seg_root,
json_file=instances_json, # TODO rename
evaluator_type="coco_panoptic_seg",
**metadata)
semantic_name = name + "_stuffonly"
DatasetCatalog.register(semantic_name,
lambda: load_sem_seg(sem_seg_root, image_root))
MetadataCatalog.get(semantic_name).set(sem_seg_root=sem_seg_root,
image_root=image_root,
evaluator_type="sem_seg",
**metadata)
def get_midpoint_obj_conf(self):
#objects = random.sample(list(objects), self.num_objects_per_episode)
xyz_obj_mids = []
count = 0
nav_points = np.array(self.nav_pts)
action = "do_nothing"
movement = "turn_right"
scene = self.sim.semantic_scene
objects = scene.objects
if self.visualize:
self.plot_navigable_points(self.nav_pts)
while count < self.num_objects_per_episode:
print("GETTING OBJECT #", count)
obj_ind = np.random.randint(low=0, high=len(objects))
obj = objects[obj_ind]
if obj == None or obj.category == None or obj.category.name(
) not in self.include_classes:
continue
# st()
if self.verbose:
print(f"Object name is: {obj.category.name()}")
# Calculate distance to object center
obj_center = obj.obb.to_aabb().center
#print(obj_center)
obj_center = np.expand_dims(obj_center, axis=0)
#print(obj_center)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where(
(distances > self.radius_min) * (distances < self.radius_max))]
rand_ind = np.random.randint(low=0, high=len(valid_pts))
valid_pts = np.array(valid_pts)
pos_rand = valid_pts[rand_ind, :]
agent_state = habitat_sim.AgentState()
agent_state.position = pos_rand + np.array([0, 1.5, 0])
print("Random POS=", agent_state.position)
self.agent.set_state(agent_state)
# if self.visualize:
# x_sample = self.nav_pts[:,0]
# z_sample = self.nav_pts[:,2]
# plt.plot(z_sample, x_sample, 'o', color = 'red')
# plt.plot(agent_state.position[2], agent_state.position[0], 'x', 'blue')
# plt.show()
# rotate in place until get high confident object
# bin_angle_size = 60.0
# angles = np.arange(-180, 180, bin_angle_size)
angles = np.arange(0, 360, self.turn_angle)
for angle in angles: #b_inds_notempty:
# print("ANGLE=", angle)
# turn_angle = np.radians(angle)
# quat_yaw = quat_from_angle_axis(angle, np.array([0, 1.0, 0]))
# # Set agent yaw rotation to current angle
# agent_state.rotation = quat_yaw
# # change sensor state to default
# # need to move the sensors too
# for sensor in self.agent.state.sensor_states:
# # st()
# self.agent.state.sensor_states[sensor].rotation = agent_state.rotation
# self.agent.state.sensor_states[sensor].position = agent_state.position # + np.array([0, 1.5, 0]) # ADDED IN UP TOP
# # print("PRINT", self.agent.state.sensor_states[sensor].rotation)
# print(agent_state.rotation)
# # get observations after centiering
# self.agent.set_state(agent_state)
# rotate until find confident object
observations = self.sim.step(movement) #self.sim.step(action)
####### %%%%%%%%%%%%%%%%%%%%%%% ######### MASK RCNN
im = observations["color_sensor"]
im = Image.fromarray(im, mode="RGBA")
im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
# plt.imshow(im)
# plt.show()
outputs = self.maskrcnn(im)
pred_masks = outputs['instances'].pred_masks
pred_boxes = outputs['instances'].pred_boxes.tensor
pred_classes = outputs['instances'].pred_classes
pred_scores = outputs['instances'].scores
maskrcnn_to_catname = {
56: "chair",
59: "bed",
61: "toilet",
57: "couch",
58: "indoor-plant",
72: "refrigerator",
62: "tv",
60: "dining-table"
}
obj_ids = []
obj_catids = []
obj_scores = []
obj_masks = []
obj_all_catids = []
obj_all_scores = []
obj_all_boxes = []
for segs in range(len(pred_masks)):
if pred_classes[segs].item() in maskrcnn_to_catname:
if pred_scores[segs] >= 0.90:
obj_ids.append(segs)
obj_catids.append(pred_classes[segs].item())
obj_scores.append(pred_scores[segs].item())
obj_masks.append(pred_masks[segs])
obj_all_catids.append(pred_classes[segs].item())
obj_all_scores.append(pred_scores[segs].item())
y, x = torch.where(pred_masks[segs])
pred_box = torch.Tensor(
[min(y), min(x),
max(y), max(x)]) # ymin, xmin, ymax, xmax
obj_all_boxes.append(pred_box)
print("MASKS ", len(pred_masks))
print("VALID ", len(obj_scores))
print(obj_scores)
print(pred_scores.shape)
translation_ = self.agent.state.sensor_states[
'depth_sensor'].position
quaternion_ = self.agent.state.sensor_states[
'depth_sensor'].rotation
rotation_ = quaternion.as_rotation_matrix(quaternion_)
T_world_cam = np.eye(4)
T_world_cam[0:3, 0:3] = rotation_
T_world_cam[0:3, 3] = translation_
if not obj_masks:
continue
else:
# randomly choose a high confidence object
# instead of this I think we should iterate over ALL the high confident objects and fixate on them
obj_mask_focus = random.choice(obj_masks)
depth = observations["depth_sensor"]
xs, ys = np.meshgrid(
np.linspace(-1 * 256 / 2., 1 * 256 / 2., 256),
np.linspace(1 * 256 / 2., -1 * 256 / 2., 256))
depth = depth.reshape(1, 256, 256)
xs = xs.reshape(1, 256, 256)
ys = ys.reshape(1, 256, 256)
xys = np.vstack(
(xs * depth, ys * depth, -depth, np.ones(depth.shape)))
xys = xys.reshape(4, -1)
xy_c0 = np.matmul(np.linalg.inv(self.K), xys)
xyz = xy_c0.T[:, :3].reshape(256, 256, 3)
xyz_obj_masked = xyz[obj_mask_focus]
xyz_obj_masked = np.matmul(
rotation_, xyz_obj_masked.T) + translation_.reshape(
3, 1)
xyz_obj_mid = np.mean(xyz_obj_masked, axis=1)
print("MIDPOINT=", xyz_obj_mid)
xyz_obj_mids.append(xyz_obj_mid)
count += 1
if self.visualize:
plt.figure(1)
v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(
self.cfg_det.DATASETS.TRAIN[0]),
scale=1.2)
out = v.draw_instance_predictions(
outputs['instances'].to("cpu"))
seg_im = out.get_image()
plt.imshow(seg_im)
plt.figure(2)
x_sample = self.nav_pts[:, 0]
z_sample = self.nav_pts[:, 2]
plt.plot(z_sample, x_sample, 'o', color='red')
plt.plot(valid_pts[:, 2],
valid_pts[:, 0],
'o',
color='green')
plt.plot(obj_center[:, 2],
obj_center[:, 0],
'x',
color='black')
plt.plot(agent_state.position[2],
agent_state.position[0],
'x',
color='blue')
plt.show()
break # got an object
xyz_obj_mids = np.array(xyz_obj_mids)
return xyz_obj_mids
def main(config):
root = expanduser(config["base"]["root"])
imgs_root = expanduser(config["base"]["imgs_root"])
jsons_dir = join(root, "jsons")
model_dir = join(root, "outputs")
scale = float(config["test_model"]["scale"])
do_show = config["test_model"]["do_show"]
register_data(jsons_dir, imgs_root)
# Need this datasets line, in order for metadata to have .thing_classes attribute
datasets = DatasetCatalog.get("test_data")
metadata = MetadataCatalog.get("test_data")
# Read the cfg back in:
with open(join(model_dir, "cfg.txt"), "r") as f:
cfg = f.read()
# Turn into CfgNode obj:
cfg = CfgNode.load_cfg(cfg)
# Use the weights from the model trained on our custom dataset:
cfg.MODEL.WEIGHTS = join(model_dir, "model_final.pth") # TODO: have option to use snapshot instead
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.01 # make small so I can make PR curve for broad range of scores
# cfg.DATASETS.TEST = ("val_data", ) # should already be saved from train_model.py
print("Generating predictor ...")
predictor = DefaultPredictor(cfg)
# For saving images with predicted labels:
output_imgs_dir = join(model_dir, "test_pred_imgs")
makedirs(output_imgs_dir, exist_ok=True)
# For saving detection predictions as csv:
output_csv = join(model_dir, "all_test_preds.csv")
csv_file_handle = open(output_csv, "w", newline="")
atexit.register(csv_file_handle.close)
col_names = ["img", "x1", "y1", "x2", "y2", "score", "thing","dummy_id"]
csv_writer = csv.DictWriter(csv_file_handle, fieldnames=col_names)
csv_writer.writeheader()
# Select 5 random images to visualize,
# but save the prediction results for all imgs:
rando_idxs = np.random.choice(range(len(datasets)), 5, replace=False).tolist()
for i,d in enumerate(datasets):
print(f"Predicting on image {i+1} of {len(datasets)} ...", end="\r")
id = d["image_id"]
img = cv2.imread(d["file_name"])
detected = predictor(img)
# Visualize:
visualizer = Visualizer(img[:, :, ::-1],
metadata=metadata,
scale=scale,
instance_mode=ColorMode)
visualizer = visualizer.draw_instance_predictions(detected["instances"].to("cpu"))
# Save the first 5 images from the random draw:
if i in rando_idxs:
pred_img = visualizer.get_image()[:, :, ::-1]
cv2.imwrite(join(output_imgs_dir, ("predicted_" + basename(d["file_name"]))), pred_img)
if do_show:
cv2.imshow(f"prediction on image {id}", pred_img)
print(f"Press any key to go to the next image ({i+1}/5) ...")
key = cv2.waitKey(0) & 0xFF
if key == ord("q"):
print("Quitting ...")
break
cv2.destroyAllWindows()
# Stream the predicted box coords and scores to a csv:
preds = detected['instances'].to('cpu')
boxes = preds.pred_boxes
thing_ids = preds.pred_classes.tolist()
scores = preds.scores
num_boxes = np.array(scores.size())[0]
for i in range(0, num_boxes):
coords = boxes[i].tensor.numpy()
score = float(scores[i].numpy())
thing_id = thing_ids[i] # is int
thing_class = metadata.thing_classes[thing_id]
csv_writer.writerow({col_names[0]: basename(d["file_name"]),
col_names[1]: int(coords[0][0]), # x1
col_names[2]: int(coords[0][1]), # y1
col_names[3]: int(coords[0][2]), # x2
col_names[4]: int(coords[0][3]), # y2
col_names[5]: score, # score
col_names[6]: thing_class, # thing
col_names[7]: i}) # dummy id
print(f"Finished predicting on all {len(datasets)} images from the test data fraction.")
print(f"Results are stored in {output_csv}")
print(f"5 sample test images are stored in {output_imgs_dir}\n"
"Note that the 5 sample test images show all detections with a score greater than 0.01. "
"This low score cutoff is for test purposes and is intentional. "
"You should expect to see many false positive labels.\n")
# Clear GPU memory
torch.cuda.empty_cache()
def metadata(self):
return MetadataCatalog.get("coco_2017_train")
