import argparse import random import numpy as np import torch from models import build_model import torchvision from torchvision.ops.boxes import batched_nms import cv2 def get_args_parser(): parser = argparse.ArgumentParser('Set transformer detector', add_help=False) parser.add_argument('--lr', default=1e-4, type=float) parser.add_argument('--lr_backbone', default=1e-5, type=float) parser.add_argument('--batch_size', default=2, type=int) parser.add_argument('--weight_decay', default=1e-4, type=float) parser.add_argument('--epochs', default=300, type=int) parser.add_argument('--lr_drop', default=200, type=int) parser.add_argument('--clip_max_norm', default=0.1, type=float, help='gradient clipping max norm') # Model parameters parser.add_argument('--frozen_weights', type=str, default=None, help="Path to the pretrained model. If set, only the mask head will be trained") # * Backbone # 如果设置为resnet101，后面的权重文件路径也需要修改一下 parser.add_argument('--backbone', default='resnet50', type=str, help="Name of the convolutional backbone to use") parser.add_argument('--dilation', action='store_true', help="If true, we replace stride with dilation in the last convolutional block (DC5)") parser.add_argument('--position_embedding', default='sine', type=str, choices=('sine', 'learned'), help="Type of positional embedding to use on top of the image features") # * Transformer parser.add_argument('--enc_layers', default=6, type=int, help="Number of encoding layers in the transformer") parser.add_argument('--dec_layers', default=6, type=int, help="Number of decoding layers in the transformer") parser.add_argument('--dim_feedforward', default=2048, type=int, help="Intermediate size of the feedforward layers in the transformer blocks") parser.add_argument('--hidden_dim', default=256, type=int, help="Size of the embeddings (dimension of the transformer)") parser.add_argument('--dropout', default=0.1, type=float, help="Dropout applied in the transformer") parser.add_argument('--nheads', default=8, type=int, help="Number of attention heads inside the transformer's attentions") parser.add_argument('--num_queries', default=100, type=int, help="Number of query slots") parser.add_argument('--pre_norm', action='store_true') # * Segmentation parser.add_argument('--masks', action='store_true', help="Train segmentation head if the flag is provided") # Loss parser.add_argument('--no_aux_loss', dest='aux_loss', default='False', help="Disables auxiliary decoding losses (loss at each layer)") # * Matcher parser.add_argument('--set_cost_class', default=1, type=float, help="Class coefficient in the matching cost") parser.add_argument('--set_cost_bbox', default=5, type=float, help="L1 box coefficient in the matching cost") parser.add_argument('--set_cost_giou', default=2, type=float, help="giou box coefficient in the matching cost") # * Loss coefficients parser.add_argument('--mask_loss_coef', default=1, type=float) parser.add_argument('--dice_loss_coef', default=1, type=float) parser.add_argument('--bbox_loss_coef', default=5, type=float) parser.add_argument('--giou_loss_coef', default=2, type=float) parser.add_argument('--eos_coef', default=0.1, type=float, help="Relative classification weight of the no-object class") # dataset parameters parser.add_argument('--dataset_file', default='coco') parser.add_argument('--coco_path', type=str, default="coco") parser.add_argument('--coco_panoptic_path', type=str) parser.add_argument('--remove_difficult', action='store_true') # 检测的图像路径 parser.add_argument('--source_dir', default='demo/images', help='path where to save, empty for no saving') # 检测结果保存路径 parser.add_argument('--output_dir', default='demo/outputs', help='path where to save, empty for no saving') parser.add_argument('--device', default='cpu', help='device to use for training / testing') parser.add_argument('--seed', default=42, type=int) # resnet50对应的权重文件 parser.add_argument('--resume', default='demo/weights/detr-r50-e632da11.pth', help='resume from checkpoint') parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch') parser.add_argument('--eval', default="True") parser.add_argument('--num_workers', default=2, type=int) # distributed training parameters parser.add_argument('--world_size', default=1, type=int, help='number of distributed processes') parser.add_argument('--dist_url', default='env://', help='url used to set up distributed training') return parser def init(): args, _ = get_args_parser().parse_known_args() device = torch.device(args.device) model, _, _ = build_model(args) Checkpoint = torch.load(args.resume, map_location="cpu") model.load_state_dict(Checkpoint["model"], False) model.to(device) return model, device # COCO classes Classes = [ 'N/A', 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella', 'N/A', 'N/A', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'N/A', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'N/A', 'dining table', 'N/A', 'N/A', 'toilet', 'N/A', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'N/A', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush' ] Model, Device = init() ToTensor = torchvision.transforms.ToTensor() def box_cxcywh_to_xyxy(x): x_c, y_c, w, h = x.unbind(1) b = [(x_c - 0.5 * w), (y_c - 0.5 * h), (x_c + 0.5 * w), (y_c + 0.5 * h)] return torch.stack(b, dim=1) def rescale_bboxes(out_bbox, size): img_w, img_h = size b = box_cxcywh_to_xyxy(out_bbox) b = b * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32) return b def filter_boxes(scores, boxes, confidence=0.7, apply_nms=True, iou=0.5): keep = scores.max(-1).values > confidence scores, boxes = scores[keep], boxes[keep] if apply_nms: top_scores, labels = scores.max(-1) keep = batched_nms(boxes, top_scores, labels, iou) scores, boxes = scores[keep], boxes[keep] return scores, boxes def plot_one_box(x, img, color=None, label=None, line_thickness=1): tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness color = color or [random.randint(0, 255) for _ in range(3)] c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3])) cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA) if label: tf = max(tl - 1, 1) # font thickness t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0] c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3 cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA) def inference(img: "np.ndarray") -> "np.ndarray": # 加载图片数据到cpu、gpu中 imgTensor = ToTensor(img) imgTensor = torch.reshape(imgTensor, [-1, imgTensor.shape[0], imgTensor.shape[1], imgTensor.shape[2]]) imgTensor.to(Device) # Model(image)即可检测图片里的对象数据 inferenceResult = Model(imgTensor) # 检测对象的得分 scores = inferenceResult["pred_logits"].softmax(-1)[0, :, :-1].cpu() # 检测对象的位置数据 boxes = rescale_bboxes(inferenceResult["pred_boxes"][0,].cpu(), (imgTensor.shape[3], imgTensor.shape[2])) scores, boxes = filter_boxes(scores, boxes) scores, boxes = scores.data.numpy(), boxes.data.numpy() # 在图片中标记对象 for i in range(boxes.shape[0]): class_id = scores[i].argmax() label = Classes[class_id] confidence = scores[i].max() text = f"{label} {confidence:.3f}" plot_one_box(boxes[i], img, label=text) # 返回标记后的图像数据 return img def main(): img = cv2.imread("./123.png") out = inference(img) cv2.imwrite("./out.jpg", out) if __name__ == "__main__": main()