yoloV3 单张图像测试（slimYoloV3）

import argparse
import json

from torch.utils.data import DataLoader

from models import *
from utils.datasets import *
from utils.utils import *

def test_one_image(     cfg,
                        data,
                        weights=None,
                        batch_size=16,
                        img_size=416,
                        iou_thres=0.5,
                        conf_thres=0.001,
                        nms_thres=0.5,
                        save_json=False,
                        model=None):
    if model is None:
        device = torch_utils.select_device(opt.device)
        verbose = True

        # Initialize model
        model = Darknet(cfg, img_size).to(device)

        # Load weights
        attempt_download(weights)
        if weights.endswith('.pt'):  # pytorch format
            model.load_state_dict(torch.load(weights, map_location=device)['model'])
        else:  # darknet format
            _ = load_darknet_weights(model, weights)

        if torch.cuda.device_count() > 1:
            model = nn.DataParallel(model)
    else:
        device = next(model.parameters()).device  # get model device
        verbose = False
    
    # Configure run
    data = parse_data_cfg(data)
    nc = int(data['classes'])  # number of classes
    test_path = data['valid']  # path to test images
    names = load_classes(data['names'])  # class names
    
    seen = 0
    model.eval()
    coco91class = coco80_to_coco91_class()
    s = ('%20s' + '%10s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP', 'F1')

    img = None
    img_size=416
    img_path = 'data/samples/bus.jpg'
    if img is None:
        img = cv2.imread(img_path)  # BGR
        assert img is not None, 'Image Not Found ' + img_path
        r = img_size / max(img.shape)  # size ratio
        # if r < 1:  # if training (NOT testing), downsize to inference shape
        h, w, _ = img.shape
        
        img = cv2.resize(img, (int(w * r), int(h * r)), interpolation=cv2.INTER_LINEAR)  # INTER_LINEAR fastest
    # Normalize
    shapes = (h, w)
    
    img_show = img.copy()
    img_show = np.pad(img_show,((0,0),(52,52),(0,0)),'constant',constant_values = (0,0))
    
    print("1------", img_show.shape)
    img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, to 3x416x416
    img = np.pad(img,((0,0),(0,0),(52,52)),'constant',constant_values = (0,0))
    
    print("2------", img_show.shape)
    img = np.expand_dims(img,axis=0)
    img = np.ascontiguousarray(img, dtype=np.float32)  # uint8 to float32
    img /= 255.0  # 0 - 255 to 0.0 - 1.0
    batch, channel, height, width = img.shape  # batch size, channels, height, width
    print("------",batch, channel, height, width)
    img_ = torch.from_numpy(img)
    inf_out, train_out = model(img_)  # inference and training outputs
    # Compute loss
    if hasattr(model, 'hyp'):  # if model has loss hyperparameters
        loss += compute_loss(train_out, targets, model)[1][:3].cpu()  # GIoU, obj, cls
    # Run NMS
    # (x1, y1, x2, y2, object_conf, class_conf, class)
    output = non_max_suppression(inf_out, conf_thres=conf_thres, nms_thres=nms_thres)
    # print("output:", output)
    # Statistics per image
    for si, pred in enumerate(output):
        box = pred[:, :4].clone()  # xyxy
        conf = pred[:, 4:].clone()  # xyxy
        # scale_coords(img.shape[1:], box, shapes)  # to original shape
        # scale_coords(shapes, box, img.shape[1:])  # to original shape
        # box = xyxy2xywh(box)  # xywh
        # box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner
        
        box = box.numpy()
        conf = conf.numpy()
        for i in range(box.shape[0]):
            b = box[i,:]
            c = conf[i,:]
            if c[0]>0.9 and c[1]>0.9:
                cv2.rectangle(img_show, (b[0], b[1]), (b[2], b[3]), (0, 255, 0), 2)
                print(b,c)
                # print(img_show)
        cv2.imwrite("1.jpg", img_show)
            
if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog='test.py')
    parser.add_argument('--cfg', type=str, default='cfg/yolov3-spp.cfg', help='cfg file path')
    parser.add_argument('--data', type=str, default='data/coco.data', help='coco.data file path')
    # yolov3-spp.ptyolov3-spp.pt
    parser.add_argument('--weights', type=str, default='weights/yolov3-spp.weights', help='path to weights file')
    parser.add_argument('--batch-size', type=int, default=16, help='size of each image batch')
    parser.add_argument('--img-size', type=int, default=416, help='inference size (pixels)')
    parser.add_argument('--iou-thres', type=float, default=0.5, help='iou threshold required to qualify as detected')
    parser.add_argument('--conf-thres', type=float, default=0.001, help='object confidence threshold')
    parser.add_argument('--nms-thres', type=float, default=0.5, help='iou threshold for non-maximum suppression')
    parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
    parser.add_argument('--device', default='0', help='device id (i.e. 0 or 0,1) or cpu')
    opt = parser.parse_args()
    print(opt)

    with torch.no_grad():
        test_one_image( opt.cfg,
                        opt.data,
                        opt.weights,
                        opt.batch_size,
                        opt.img_size,
                        opt.iou_thres,
                        opt.conf_thres,
                        opt.nms_thres,
                        opt.save_json)