centernet2训练自己的数据集
环境配置
1.下载https://github.com/xingyizhou/CenterNet2
2.安装环境
conda create -n CenterNet2 python=3.8
conda activate CenterNet2
pip install torch==1.8
python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'
pip install cython opencv-python pillow matplotlib termcolor cloudpickle tabulate tensorboard termcolor tqdm yacs mock fvcore pydot wheel future omegaconf==2.1.0.dev22
文件准备
modified: detectron2/engine/defaults.py
parser.add_argument("--config-file", default="./configs/CenterNet2_R50_1x.yaml", metavar="FILE", help="path to config file")
modified: projects/CenterNet2/centernet/config.py
_C.MODEL.CENTERNET.NUM_CLASSES = 10
modified: projects/CenterNet2/configs/Base-CenterNet2.yaml
# WEIGHTS: "detectron2://ImageNetPretrained/MSRA/R-50.pkl"
WEIGHTS: "models/CenterNet2_10.pth"
BASE_LR: 0.01
STEPS: (60000, 80000)
MAX_ITER: 100
CHECKPOINT_PERIOD: 50
WARMUP_ITERS: 40
modified: projects/CenterNet2/demo.py
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
import argparse
import glob
import multiprocessing as mp
import os
import time
import cv2
import tqdm
from detectron2.config import get_cfg
from detectron2.data.detection_utils import read_image
from detectron2.utils.logger import setup_logger
from predictor import VisualizationDemo
from centernet.config import add_centernet_config
# constants
WINDOW_NAME = "CenterNet2 detections"
from detectron2.utils.video_visualizer import VideoVisualizer
from detectron2.utils.visualizer import ColorMode, Visualizer
from detectron2.data import MetadataCatalog
NUM_CLASSES=10
def setup_cfg(args):
# load config from file and command-line arguments
cfg = get_cfg()
add_centernet_config(cfg)
cfg.MODEL.CENTERNET.NUM_CLASSES=NUM_CLASSES
cfg.MODEL.ROI_HEADS.NUM_CLASSES =NUM_CLASSES
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
# Set score_threshold for builtin models
cfg.MODEL.RETINANET.SCORE_THRESH_TEST = args.confidence_threshold
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = args.confidence_threshold
if cfg.MODEL.META_ARCHITECTURE in ['ProposalNetwork', 'CenterNetDetector']:
cfg.MODEL.CENTERNET.INFERENCE_TH = args.confidence_threshold
cfg.MODEL.CENTERNET.NMS_TH = cfg.MODEL.ROI_HEADS.NMS_THRESH_TEST
cfg.MODEL.PANOPTIC_FPN.COMBINE.INSTANCES_CONFIDENCE_THRESH = args.confidence_threshold
cfg.freeze()
return cfg
def get_parser():
parser = argparse.ArgumentParser(description="Detectron2 demo for builtin models")
parser.add_argument(
"--config-file",
default="configs/CenterNet2_R50_1x.yaml",
metavar="FILE",
help="path to config file",
)
parser.add_argument("--webcam", action="store_true", help="Take inputs from webcam.")
parser.add_argument("--video-input", help="Path to video file.")
parser.add_argument("--input", default="imgs/",nargs="+", help="A list of space separated input images")
parser.add_argument(
"--output",
default="output",
help="A file or directory to save output visualizations. "
"If not given, will show output in an OpenCV window.",
)
parser.add_argument(
"--confidence-threshold",
type=float,
default=0.3,
help="Minimum score for instance predictions to be shown",
)
parser.add_argument(
"--opts",
help="Modify config options using the command-line 'KEY VALUE' pairs",
default=['MODEL.WEIGHTS', 'output/CenterNet2/CenterNet2_R50_1x/model_final.pth'],
nargs=argparse.REMAINDER,
)
return parser
if __name__ == "__main__":
mp.set_start_method("spawn", force=True)
args = get_parser().parse_args()
logger = setup_logger()
logger.info("Arguments: " + str(args))
cfg = setup_cfg(args)
demo = VisualizationDemo(cfg)
output_file = None
if args.input:
if len(args.input) == 1:
args.input = glob.glob(os.path.expanduser(args.input[0]))
files = os.listdir(args.input[0])
args.input = [args.input[0] + x for x in files]
assert args.input, "The input path(s) was not found"
visualizer = VideoVisualizer(
MetadataCatalog.get(
cfg.DATASETS.TEST[0] if len(cfg.DATASETS.TEST) else "__unused"
),
instance_mode=ColorMode.IMAGE)
for path in tqdm.tqdm(os.listdir(args.input)):
path = os.path.join(args.input, path)
# use PIL, to be consistent with evaluation
img = read_image(path, format="BGR")
start_time = time.time()
visualizer.metadata.thing_classes[:10]=["UxingquejiaoOrbengbian","Vxingquejiao",
"barcode","chuanjuanju","cuowei",
"jiaodai","liepian","lubai","pianjianju","yiwu"]
predictions, visualized_output = demo.run_on_image(
img, visualizer=visualizer)
if 'instances' in predictions:
logger.info(
"{}: detected {} instances in {:.2f}s".format(
path, len(predictions["instances"]), time.time() - start_time
)
)
else:
logger.info(
"{}: detected {} instances in {:.2f}s".format(
path, len(predictions["proposals"]), time.time() - start_time
)
)
if args.output:
if os.path.isdir(args.output):
assert os.path.isdir(args.output), args.output
out_filename = os.path.join(args.output, os.path.basename(path))
visualized_output.save(out_filename)
else:
# assert len(args.input) == 1, "Please specify a directory with args.output"
# out_filename = args.output
if output_file is None:
width = visualized_output.get_image().shape[1]
height = visualized_output.get_image().shape[0]
frames_per_second = 15
output_file = cv2.VideoWriter(
filename=args.output,
# some installation of opencv may not support x264 (due to its license),
# you can try other format (e.g. MPEG)
fourcc=cv2.VideoWriter_fourcc(*"x264"),
fps=float(frames_per_second),
frameSize=(width, height),
isColor=True,
)
output_file.write(visualized_output.get_image()[:, :, ::-1])
else:
# cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
cv2.imshow(WINDOW_NAME, visualized_output.get_image()[:, :, ::-1])
if cv2.waitKey(1 ) == 27:
break # esc to quit
elif args.webcam:
assert args.input is None, "Cannot have both --input and --webcam!"
cam = cv2.VideoCapture(0)
for vis in tqdm.tqdm(demo.run_on_video(cam)):
cv2.namedWindow(WINDOW_NAME, cv2.WINDOW_NORMAL)
cv2.imshow(WINDOW_NAME, vis)
if cv2.waitKey(1) == 27:
break # esc to quit
cv2.destroyAllWindows()
elif args.video_input:
video = cv2.VideoCapture(args.video_input)
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
frames_per_second = 15 # video.get(cv2.CAP_PROP_FPS)
num_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
basename = os.path.basename(args.video_input)
if args.output:
if os.path.isdir(args.output):
output_fname = os.path.join(args.output, basename)
output_fname = os.path.splitext(output_fname)[0] + ".mkv"
else:
output_fname = args.output
# assert not os.path.isfile(output_fname), output_fname
output_file = cv2.VideoWriter(
filename=output_fname,
# some installation of opencv may not support x264 (due to its license),
# you can try other format (e.g. MPEG)
fourcc=cv2.VideoWriter_fourcc(*"x264"),
fps=float(frames_per_second),
frameSize=(width, height),
isColor=True,
)
assert os.path.isfile(args.video_input)
for vis_frame in tqdm.tqdm(demo.run_on_video(video), total=num_frames):
if args.output:
output_file.write(vis_frame)
cv2.namedWindow(basename, cv2.WINDOW_NORMAL)
cv2.imshow(basename, vis_frame)
if cv2.waitKey(1) == 27:
break # esc to quit
video.release()
if args.output:
output_file.release()
else:
cv2.destroyAllWindows()
modified: projects/CenterNet2/train_net.py
import logging
import os
from collections import OrderedDict
import torch
from torch.nn.parallel import DistributedDataParallel
import time
import datetime
import json
from fvcore.common.timer import Timer
import detectron2.utils.comm as comm
from detectron2.checkpoint import DetectionCheckpointer, PeriodicCheckpointer
from detectron2.config import get_cfg
from detectron2.data import (
MetadataCatalog,
build_detection_test_loader,
)
from detectron2.engine import default_argument_parser, default_setup, launch
from detectron2.evaluation import (
COCOEvaluator,
LVISEvaluator,
inference_on_dataset,
print_csv_format,
)
from detectron2.modeling import build_model
from detectron2.solver import build_lr_scheduler, build_optimizer
from detectron2.utils.events import (
CommonMetricPrinter,
EventStorage,
JSONWriter,
TensorboardXWriter,
)
from detectron2.modeling.test_time_augmentation import GeneralizedRCNNWithTTA
from detectron2.data.dataset_mapper import DatasetMapper
from detectron2.data.build import build_detection_train_loader
from centernet.config import add_centernet_config
from centernet.data.custom_build_augmentation import build_custom_augmentation
from detectron2.data.datasets import register_coco_instances
logger = logging.getLogger("detectron2")
def do_test(cfg, model):
results = OrderedDict()
for dataset_name in cfg.DATASETS.TEST:
mapper = None if cfg.INPUT.TEST_INPUT_TYPE == 'default' else \
DatasetMapper(
cfg, False, augmentations=build_custom_augmentation(cfg, False))
data_loader = build_detection_test_loader(cfg, dataset_name, mapper=mapper)
output_folder = os.path.join(
cfg.OUTPUT_DIR, "inference_{}".format(dataset_name))
evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
if evaluator_type == "lvis":
evaluator = LVISEvaluator(dataset_name, cfg, True, output_folder)
elif evaluator_type == 'coco':
evaluator = COCOEvaluator(dataset_name, cfg, True, output_folder)
else:
assert 0, evaluator_type
results[dataset_name] = inference_on_dataset(
model, data_loader, evaluator)
if comm.is_main_process():
logger.info("Evaluation results for {} in csv format:".format(
dataset_name))
print_csv_format(results[dataset_name])
if len(results) == 1:
results = list(results.values())[0]
return results
def do_train(cfg, model, resume=False):
model.train()
optimizer = build_optimizer(cfg, model)
scheduler = build_lr_scheduler(cfg, optimizer)
checkpointer = DetectionCheckpointer(
model, cfg.OUTPUT_DIR, optimizer=optimizer, scheduler=scheduler
)
start_iter = (
checkpointer.resume_or_load(
cfg.MODEL.WEIGHTS, resume=resume,
).get("iteration", -1) + 1
)
if cfg.SOLVER.RESET_ITER:
logger.info('Reset loaded iteration. Start training from iteration 0.')
start_iter = 0
max_iter = cfg.SOLVER.MAX_ITER if cfg.SOLVER.TRAIN_ITER < 0 else cfg.SOLVER.TRAIN_ITER
periodic_checkpointer = PeriodicCheckpointer(
checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD, max_iter=max_iter
)
writers = (
[
CommonMetricPrinter(max_iter),
JSONWriter(os.path.join(cfg.OUTPUT_DIR, "metrics.json")),
TensorboardXWriter(cfg.OUTPUT_DIR),
]
if comm.is_main_process()
else []
)
mapper = DatasetMapper(cfg, True) if cfg.INPUT.CUSTOM_AUG == '' else \
DatasetMapper(cfg, True, augmentations=build_custom_augmentation(cfg, True))
if cfg.DATALOADER.SAMPLER_TRAIN in ['TrainingSampler', 'RepeatFactorTrainingSampler']:
data_loader = build_detection_train_loader(cfg, mapper=mapper)
else:
from centernet.data.custom_dataset_dataloader import build_custom_train_loader
data_loader = build_custom_train_loader(cfg, mapper=mapper)
logger.info("Starting training from iteration {}".format(start_iter))
with EventStorage(start_iter) as storage:
step_timer = Timer()
data_timer = Timer()
start_time = time.perf_counter()
for data, iteration in zip(data_loader, range(start_iter, max_iter)):
data_time = data_timer.seconds()
storage.put_scalars(data_time=data_time)
step_timer.reset()
iteration = iteration + 1
storage.step()
loss_dict = model(data)
losses = sum(
loss for k, loss in loss_dict.items())
assert torch.isfinite(losses).all(), loss_dict
loss_dict_reduced = {k: v.item() \
for k, v in comm.reduce_dict(loss_dict).items()}
losses_reduced = sum(loss for loss in loss_dict_reduced.values())
if comm.is_main_process():
storage.put_scalars(
total_loss=losses_reduced, **loss_dict_reduced)
optimizer.zero_grad()
losses.backward()
optimizer.step()
storage.put_scalar(
"lr", optimizer.param_groups[0]["lr"], smoothing_hint=False)
step_time = step_timer.seconds()
storage.put_scalars(time=step_time)
data_timer.reset()
scheduler.step()
if (
cfg.TEST.EVAL_PERIOD > 0
and iteration % cfg.TEST.EVAL_PERIOD == 0
and iteration != max_iter
):
do_test(cfg, model)
comm.synchronize()
if iteration - start_iter > 5 and \
(iteration % 20 == 0 or iteration == max_iter):
for writer in writers:
writer.write()
periodic_checkpointer.step(iteration)
total_time = time.perf_counter() - start_time
logger.info(
"Total training time: {}".format(
str(datetime.timedelta(seconds=int(total_time)))))
NUM_CLASSES=10
def setup(args):
"""
Create configs and perform basic setups.
"""
register_coco_instances("train", {}, "datasets/coco/annotations/instances_train2017.json", "datasets/coco/train2017")
register_coco_instances("test", {}, "datasets/coco/annotations/instances_val2017.json", "datasets/coco/val2017")
cfg = get_cfg()
add_centernet_config(cfg)
cfg.merge_from_file(args.config_file)
cfg.merge_from_list(args.opts)
cfg.DATASETS.TRAIN = ("train",)
cfg.DATASETS.TEST = ("test",)
cfg.MODEL.CENTERNET.NUM_CLASSES=NUM_CLASSES
cfg.MODEL.ROI_HEADS.NUM_CLASSES = NUM_CLASSES
if '/auto' in cfg.OUTPUT_DIR:
file_name = os.path.basename(args.config_file)[:-5]
cfg.OUTPUT_DIR = cfg.OUTPUT_DIR.replace('/auto', '/{}'.format(file_name))
logger.info('OUTPUT_DIR: {}'.format(cfg.OUTPUT_DIR))
cfg.freeze()
default_setup(cfg, args)
return cfg
def main(args):
cfg = setup(args)
model = build_model(cfg)
logger.info("Model:\n{}".format(model))
if args.eval_only:
DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load(
cfg.MODEL.WEIGHTS, resume=args.resume
)
if cfg.TEST.AUG.ENABLED:
logger.info("Running inference with test-time augmentation ...")
model = GeneralizedRCNNWithTTA(cfg, model, batch_size=1)
return do_test(cfg, model)
distributed = comm.get_world_size() > 1
if distributed:
model = DistributedDataParallel(
model, device_ids=[comm.get_local_rank()], broadcast_buffers=False,
find_unused_parameters=True
)
do_train(cfg, model, resume=args.resume)
return do_test(cfg, model)
if __name__ == "__main__":
args = default_argument_parser()
args.add_argument('--manual_device', default='')
args = args.parse_args()
if args.manual_device != '':
os.environ['CUDA_VISIBLE_DEVICES'] = args.manual_device
args.dist_url = 'tcp://127.0.0.1:{}'.format(
torch.randint(11111, 60000, (1,))[0].item())
print("Command Line Args:", args)
launch(
main,
args.num_gpus,
num_machines=args.num_machines,
machine_rank=args.machine_rank,
dist_url=args.dist_url,
args=(args,),
)
3.新建文件:
CenterNet2\projects\CenterNet2\changecenternet2.py
import torch
import numpy as np
import pickle
num_class = 10
pretrained_weights = torch.load('models/CenterNet2_R50_1x.pth')
pretrained_weights['iteration']=0
pretrained_weights['model']["roi_heads.box_predictor.0.cls_score.weight"].resize_(num_class+1,1024)
pretrained_weights['model']["roi_heads.box_predictor.0.cls_score.bias"].resize_(num_class+1)
pretrained_weights['model']["roi_heads.box_predictor.1.cls_score.weight"].resize_(num_class+1,1024)
pretrained_weights['model']["roi_heads.box_predictor.1.cls_score.bias"].resize_(num_class+1)
pretrained_weights['model']["roi_heads.box_predictor.2.cls_score.weight"].resize_(num_class+1,1024)
pretrained_weights['model']["roi_heads.box_predictor.2.cls_score.bias"].resize_(num_class+1)
torch.save(pretrained_weights, "models/CenterNet2_%d.pth"%num_class)
数据集准备
使用labelme标注好后,将标注好的coco数据集放到:CenterNet2/projects/CenterNet2/datasets/coco
annotations:
train2017和val2017:
json转coco的脚本请看:https://github.com/carlsummer/python_developer_tools/blob/main/python_developer_tools/cv/datasets/labelme2coco.py
开始训练和预测
python projects/CenterNet2/train_net.py
python projects/CenterNet2/demo.py