Abhilashvj/planogram-compliance

1# YOLOv5 🚀 by Ultralytics, GPL-3.0 license2"""3Validate a trained YOLOv5 model accuracy on a custom dataset4 5Usage:6    $ python path/to/val.py --data coco128.yaml --weights yolov5s.pt --img 6407"""8 9import argparse10import json11import os12import sys13from pathlib import Path14from threading import Thread15 16import numpy as np17import torch18from tqdm import tqdm19 20FILE = Path(__file__).absolute()21sys.path.append(FILE.parents[0].as_posix())  # add yolov5/ to path22 23from models.experimental import attempt_load24from utils.callbacks import Callbacks25from utils.datasets import create_dataloader26from utils.general import (27    box_iou,28    check_dataset,29    check_img_size,30    check_requirements,31    check_suffix,32    check_yaml,33    coco80_to_coco91_class,34    colorstr,35    increment_path,36    non_max_suppression,37    scale_coords,38    set_logging,39    xywh2xyxy,40    xyxy2xywh,41)42from utils.metrics import ConfusionMatrix, ap_per_class43from utils.plots import output_to_target, plot_images, plot_study_txt44from utils.torch_utils import select_device, time_sync45 46 47def save_one_txt(predn, save_conf, shape, file):48    # Save one txt result49    gn = torch.tensor(shape)[[1, 0, 1, 0]]  # normalization gain whwh50    for *xyxy, conf, cls in predn.tolist():51        xywh = (52            (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()53        )  # normalized xywh54        line = (55            (cls, *xywh, conf) if save_conf else (cls, *xywh)56        )  # label format57        with open(file, "a") as f:58            f.write(("%g " * len(line)).rstrip() % line + "\n")59 60 61def save_one_json(predn, jdict, path, class_map):62    # Save one JSON result {"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}63    image_id = int(path.stem) if path.stem.isnumeric() else path.stem64    box = xyxy2xywh(predn[:, :4])  # xywh65    box[:, :2] -= box[:, 2:] / 2  # xy center to top-left corner66    for p, b in zip(predn.tolist(), box.tolist()):67        jdict.append(68            {69                "image_id": image_id,70                "category_id": class_map[int(p[5])],71                "bbox": [round(x, 3) for x in b],72                "score": round(p[4], 5),73            }74        )75 76 77def process_batch(detections, labels, iouv):78    """79    Return correct predictions matrix. Both sets of boxes are in (x1, y1, x2, y2) format.80    Arguments:81        detections (Array[N, 6]), x1, y1, x2, y2, conf, class82        labels (Array[M, 5]), class, x1, y1, x2, y283    Returns:84        correct (Array[N, 10]), for 10 IoU levels85    """86    correct = torch.zeros(87        detections.shape[0],88        iouv.shape[0],89        dtype=torch.bool,90        device=iouv.device,91    )92    iou = box_iou(labels[:, 1:], detections[:, :4])93    x = torch.where(94        (iou >= iouv[0]) & (labels[:, 0:1] == detections[:, 5])95    )  # IoU above threshold and classes match96    if x[0].shape[0]:97        matches = (98            torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1)99            .cpu()100            .numpy()101        )  # [label, detection, iou]102        if x[0].shape[0] > 1:103            matches = matches[matches[:, 2].argsort()[::-1]]104            matches = matches[np.unique(matches[:, 1], return_index=True)[1]]105            # matches = matches[matches[:, 2].argsort()[::-1]]106            matches = matches[np.unique(matches[:, 0], return_index=True)[1]]107        matches = torch.Tensor(matches).to(iouv.device)108        correct[matches[:, 1].long()] = matches[:, 2:3] >= iouv109    return correct110 111 112@torch.no_grad()113def run(114    data,115    weights=None,  # model.pt path(s)116    batch_size=32,  # batch size117    imgsz=640,  # inference size (pixels)118    conf_thres=0.001,  # confidence threshold119    iou_thres=0.6,  # NMS IoU threshold120    task="val",  # train, val, test, speed or study121    device="",  # cuda device, i.e. 0 or 0,1,2,3 or cpu122    single_cls=False,  # treat as single-class dataset123    augment=False,  # augmented inference124    verbose=False,  # verbose output125    save_txt=False,  # save results to *.txt126    save_hybrid=False,  # save label+prediction hybrid results to *.txt127    save_conf=False,  # save confidences in --save-txt labels128    save_json=False,  # save a COCO-JSON results file129    project="runs/val",  # save to project/name130    name="exp",  # save to project/name131    exist_ok=False,  # existing project/name ok, do not increment132    half=True,  # use FP16 half-precision inference133    model=None,134    dataloader=None,135    save_dir=Path(""),136    plots=True,137    callbacks=Callbacks(),138    compute_loss=None,139):140    # Initialize/load model and set device141    training = model is not None142    if training:  # called by train.py143        device = next(model.parameters()).device  # get model device144 145    else:  # called directly146        device = select_device(device, batch_size=batch_size)147 148        # Directories149        save_dir = increment_path(150            Path(project) / name, exist_ok=exist_ok151        )  # increment run152        (save_dir / "labels" if save_txt else save_dir).mkdir(153            parents=True, exist_ok=True154        )  # make dir155 156        # Load model157        check_suffix(weights, ".pt")158        model = attempt_load(weights, map_location=device)  # load FP32 model159        gs = max(int(model.stride.max()), 32)  # grid size (max stride)160        imgsz = check_img_size(imgsz, s=gs)  # check image size161 162        # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99163        # if device.type != 'cpu' and torch.cuda.device_count() > 1:164        #     model = nn.DataParallel(model)165 166        # Data167        data = check_dataset(data)  # check168 169    # Half170    half &= device.type != "cpu"  # half precision only supported on CUDA171    if half:172        model.half()173 174    # Configure175    model.eval()176    is_coco = isinstance(data.get("val"), str) and data["val"].endswith(177        "coco/val2017.txt"178    )  # COCO dataset179    nc = 1 if single_cls else int(data["nc"])  # number of classes180    iouv = torch.linspace(0.5, 0.95, 10).to(181        device182    )  # iou vector for mAP@0.5:0.95183    niou = iouv.numel()184 185    # Dataloader186    if not training:187        if device.type != "cpu":188            model(189                torch.zeros(1, 3, imgsz, imgsz)190                .to(device)191                .type_as(next(model.parameters()))192            )  # run once193        task = (194            task if task in ("train", "val", "test") else "val"195        )  # path to train/val/test images196        dataloader = create_dataloader(197            data[task],198            imgsz,199            batch_size,200            gs,201            single_cls,202            pad=0.5,203            rect=True,204            prefix=colorstr(f"{task}: "),205        )[0]206 207    seen = 0208    confusion_matrix = ConfusionMatrix(nc=nc)209    names = {210        k: v211        for k, v in enumerate(212            model.names if hasattr(model, "names") else model.module.names213        )214    }215    class_map = coco80_to_coco91_class() if is_coco else list(range(1000))216    s = ("%20s" + "%11s" * 6) % (217        "Class",218        "Images",219        "Labels",220        "P",221        "R",222        "mAP@.5",223        "mAP@.5:.95",224    )225    dt, p, r, f1, mp, mr, map50, map = (226        [0.0, 0.0, 0.0],227        0.0,228        0.0,229        0.0,230        0.0,231        0.0,232        0.0,233        0.0,234    )235    loss = torch.zeros(3, device=device)236    jdict, stats, ap, ap_class = [], [], [], []237    for batch_i, (img, targets, paths, shapes) in enumerate(238        tqdm(dataloader, desc=s)239    ):240        t1 = time_sync()241        img = img.to(device, non_blocking=True)242        img = img.half() if half else img.float()  # uint8 to fp16/32243        img /= 255.0  # 0 - 255 to 0.0 - 1.0244        targets = targets.to(device)245        nb, _, height, width = img.shape  # batch size, channels, height, width246        t2 = time_sync()247        dt[0] += t2 - t1248 249        # Run model250        out, train_out = model(251            img, augment=augment252        )  # inference and training outputs253        dt[1] += time_sync() - t2254 255        # Compute loss256        if compute_loss:257            loss += compute_loss([x.float() for x in train_out], targets)[258                1259            ]  # box, obj, cls260 261        # Run NMS262        targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(263            device264        )  # to pixels265        lb = (266            [targets[targets[:, 0] == i, 1:] for i in range(nb)]267            if save_hybrid268            else []269        )  # for autolabelling270        t3 = time_sync()271        out = non_max_suppression(272            out,273            conf_thres,274            iou_thres,275            labels=lb,276            multi_label=True,277            agnostic=single_cls,278        )279        dt[2] += time_sync() - t3280 281        # Statistics per image282        for si, pred in enumerate(out):283            labels = targets[targets[:, 0] == si, 1:]284            nl = len(labels)285            tcls = labels[:, 0].tolist() if nl else []  # target class286            path, shape = Path(paths[si]), shapes[si][0]287            seen += 1288 289            if len(pred) == 0:290                if nl:291                    stats.append(292                        (293                            torch.zeros(0, niou, dtype=torch.bool),294                            torch.Tensor(),295                            torch.Tensor(),296                            tcls,297                        )298                    )299                continue300 301            # Predictions302            if single_cls:303                pred[:, 5] = 0304            predn = pred.clone()305            scale_coords(306                img[si].shape[1:], predn[:, :4], shape, shapes[si][1]307            )  # native-space pred308 309            # Evaluate310            if nl:311                tbox = xywh2xyxy(labels[:, 1:5])  # target boxes312                scale_coords(313                    img[si].shape[1:], tbox, shape, shapes[si][1]314                )  # native-space labels315                labelsn = torch.cat(316                    (labels[:, 0:1], tbox), 1317                )  # native-space labels318                correct = process_batch(predn, labelsn, iouv)319                if plots:320                    confusion_matrix.process_batch(predn, labelsn)321            else:322                correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool)323            stats.append(324                (correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls)325            )  # (correct, conf, pcls, tcls)326 327            # Save/log328            if save_txt:329                save_one_txt(330                    predn,331                    save_conf,332                    shape,333                    file=save_dir / "labels" / (path.stem + ".txt"),334                )335            if save_json:336                save_one_json(337                    predn, jdict, path, class_map338                )  # append to COCO-JSON dictionary339            callbacks.run(340                "on_val_image_end", pred, predn, path, names, img[si]341            )342 343        # Plot images344        if plots and batch_i < 3:345            f = save_dir / f"val_batch{batch_i}_labels.jpg"  # labels346            Thread(347                target=plot_images,348                args=(img, targets, paths, f, names),349                daemon=True,350            ).start()351            f = save_dir / f"val_batch{batch_i}_pred.jpg"  # predictions352            Thread(353                target=plot_images,354                args=(img, output_to_target(out), paths, f, names),355                daemon=True,356            ).start()357 358    # Compute statistics359    stats = [np.concatenate(x, 0) for x in zip(*stats)]  # to numpy360    if len(stats) and stats[0].any():361        p, r, ap, f1, ap_class = ap_per_class(362            *stats, plot=plots, save_dir=save_dir, names=names363        )364        ap50, ap = ap[:, 0], ap.mean(1)  # AP@0.5, AP@0.5:0.95365        mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean()366        nt = np.bincount(367            stats[3].astype(np.int64), minlength=nc368        )  # number of targets per class369    else:370        nt = torch.zeros(1)371 372    # Print results373    pf = "%20s" + "%11i" * 2 + "%11.3g" * 4  # print format374    print(pf % ("all", seen, nt.sum(), mp, mr, map50, map))375 376    # Print results per class377    if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats):378        for i, c in enumerate(ap_class):379            print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i]))380 381    # Print speeds382    t = tuple(x / seen * 1e3 for x in dt)  # speeds per image383    if not training:384        shape = (batch_size, 3, imgsz, imgsz)385        print(386            f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}"387            % t388        )389 390    # Plots391    if plots:392        confusion_matrix.plot(save_dir=save_dir, names=list(names.values()))393        callbacks.run("on_val_end")394 395    # Save JSON396    if save_json and len(jdict):397        w = (398            Path(weights[0] if isinstance(weights, list) else weights).stem399            if weights is not None400            else ""401        )  # weights402        anno_json = str(403            Path(data.get("path", "../coco"))404            / "annotations/instances_val2017.json"405        )  # annotations json406        pred_json = str(save_dir / f"{w}_predictions.json")  # predictions json407        print(f"\nEvaluating pycocotools mAP... saving {pred_json}...")408        with open(pred_json, "w") as f:409            json.dump(jdict, f)410 411        try:  # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb412            check_requirements(["pycocotools"])413            from pycocotools.coco import COCO414            from pycocotools.cocoeval import COCOeval415 416            anno = COCO(anno_json)  # init annotations api417            pred = anno.loadRes(pred_json)  # init predictions api418            eval = COCOeval(anno, pred, "bbox")419            if is_coco:420                eval.params.imgIds = [421                    int(Path(x).stem) for x in dataloader.dataset.img_files422                ]  # image IDs to evaluate423            eval.evaluate()424            eval.accumulate()425            eval.summarize()426            map, map50 = eval.stats[427                :2428            ]  # update results (mAP@0.5:0.95, mAP@0.5)429        except Exception as e:430            print(f"pycocotools unable to run: {e}")431 432    # Return results433    model.float()  # for training434    if not training:435        s = (436            f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}"437            if save_txt438            else ""439        )440        print(f"Results saved to {colorstr('bold', save_dir)}{s}")441    maps = np.zeros(nc) + map442    for i, c in enumerate(ap_class):443        maps[c] = ap[i]444    return (445        (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()),446        maps,447        t,448    )449 450 451def parse_opt():452    parser = argparse.ArgumentParser(prog="val.py")453    parser.add_argument(454        "--data",455        type=str,456        default="data/coco128.yaml",457        help="dataset.yaml path",458    )459    parser.add_argument(460        "--weights",461        nargs="+",462        type=str,463        default="yolov5s.pt",464        help="model.pt path(s)",465    )466    parser.add_argument(467        "--batch-size", type=int, default=32, help="batch size"468    )469    parser.add_argument(470        "--imgsz",471        "--img",472        "--img-size",473        type=int,474        default=640,475        help="inference size (pixels)",476    )477    parser.add_argument(478        "--conf-thres", type=float, default=0.001, help="confidence threshold"479    )480    parser.add_argument(481        "--iou-thres", type=float, default=0.6, help="NMS IoU threshold"482    )483    parser.add_argument(484        "--task", default="val", help="train, val, test, speed or study"485    )486    parser.add_argument(487        "--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu"488    )489    parser.add_argument(490        "--single-cls",491        action="store_true",492        help="treat as single-class dataset",493    )494    parser.add_argument(495        "--augment", action="store_true", help="augmented inference"496    )497    parser.add_argument(498        "--verbose", action="store_true", help="report mAP by class"499    )500    parser.add_argument(501        "--save-txt", action="store_true", help="save results to *.txt"502    )503    parser.add_argument(504        "--save-hybrid",505        action="store_true",506        help="save label+prediction hybrid results to *.txt",507    )508    parser.add_argument(509        "--save-conf",510        action="store_true",511        help="save confidences in --save-txt labels",512    )513    parser.add_argument(514        "--save-json",515        action="store_true",516        help="save a COCO-JSON results file",517    )518    parser.add_argument(519        "--project", default="runs/val", help="save to project/name"520    )521    parser.add_argument("--name", default="exp", help="save to project/name")522    parser.add_argument(523        "--exist-ok",524        action="store_true",525        help="existing project/name ok, do not increment",526    )527    parser.add_argument(528        "--half", action="store_true", help="use FP16 half-precision inference"529    )530    opt = parser.parse_args()531    opt.save_json |= opt.data.endswith("coco.yaml")532    opt.save_txt |= opt.save_hybrid533    opt.data = check_yaml(opt.data)  # check YAML534    return opt535 536 537def main(opt):538    set_logging()539    print(540        colorstr("val: ") + ", ".join(f"{k}={v}" for k, v in vars(opt).items())541    )542    check_requirements(543        requirements=FILE.parent / "requirements.txt",544        exclude=("tensorboard", "thop"),545    )546 547    if opt.task in ("train", "val", "test"):  # run normally548        run(**vars(opt))549 550    elif opt.task == "speed":  # speed benchmarks551        for w in (552            opt.weights if isinstance(opt.weights, list) else [opt.weights]553        ):554            run(555                opt.data,556                weights=w,557                batch_size=opt.batch_size,558                imgsz=opt.imgsz,559                conf_thres=0.25,560                iou_thres=0.45,561                save_json=False,562                plots=False,563            )564 565    elif opt.task == "study":  # run over a range of settings and save/plot566        # python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5s.pt yolov5m.pt yolov5l.pt yolov5x.pt567        x = list(range(256, 1536 + 128, 128))  # x axis (image sizes)568        for w in (569            opt.weights if isinstance(opt.weights, list) else [opt.weights]570        ):571            f = f"study_{Path(opt.data).stem}_{Path(w).stem}.txt"  # filename to save to572            y = []  # y axis573            for i in x:  # img-size574                print(f"\nRunning {f} point {i}...")575                r, _, t = run(576                    opt.data,577                    weights=w,578                    batch_size=opt.batch_size,579                    imgsz=i,580                    conf_thres=opt.conf_thres,581                    iou_thres=opt.iou_thres,582                    save_json=opt.save_json,583                    plots=False,584                )585                y.append(r + t)  # results and times586            np.savetxt(f, y, fmt="%10.4g")  # save587        os.system("zip -r study.zip study_*.txt")588        plot_study_txt(x=x)  # plot589 590 591if __name__ == "__main__":592    opt = parse_opt()593    main(opt)594