detect_image.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "b2015eab-2fa4-4fcf-af50-6842ff214433",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import torch\n",
    "from numpy import random\n",
    "\n",
    "from models.experimental import attempt_load\n",
    "from utils.datasets import LoadImages\n",
    "from utils.general import check_img_size, non_max_suppression, scale_coords\n",
    "from utils.plots import plot_one_box\n",
    "from utils.torch_utils import select_device, time_synchronized, TracedModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "280fe2fe-1536-464a-a8ee-5dcbd87b2292",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fusing layers... \n",
      "RepConv.fuse_repvgg_block\n",
      "RepConv.fuse_repvgg_block\n",
      "RepConv.fuse_repvgg_block\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\anaconda3\\envs\\pytorch-1-12\\lib\\site-packages\\torch\\_tensor.py:1083: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at  C:\\cb\\pytorch_1000000000000\\work\\build\\aten\\src\\ATen/core/TensorBody.h:482.)\n",
      "  return self._grad\n"
     ]
    }
   ],
   "source": [
    "weights = 'checkpoints/yolov7.pt'  # model.pt path(s)\n",
    "trace = False  # trace model\n",
    "\n",
    "# Initialize\n",
    "device = select_device('0')  # cuda device, i.e. 0 or 0,1,2,3 or cpu\n",
    "half = device.type != 'cpu'  # half precision only supported on CUDA\n",
    "\n",
    "# Load model\n",
    "model = attempt_load(weights, map_location=device).eval()  # load FP32 model\n",
    "stride = int(model.stride.max())  # model stride\n",
    "\n",
    "if trace:\n",
    "    model = TracedModel(model, device, img_size)\n",
    "if half:\n",
    "    model.half()  # to FP16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "d0391b3d-a584-4bbb-8bca-1d8f1abfc14c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def detect(source, view_img, save_img):\n",
    "    imgsz = check_img_size(img_size, s=stride)  # check img_size\n",
    "\n",
    "    # Set Dataloader\n",
    "    dataset = LoadImages(source, img_size=imgsz, stride=stride)\n",
    "\n",
    "    # Get names and colors\n",
    "    names = model.module.names if hasattr(model, 'module') else model.names\n",
    "    #random.seed(0)\n",
    "    colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]\n",
    "\n",
    "    # Run inference\n",
    "    if device.type != 'cpu':\n",
    "        model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters())))  # run once\n",
    "    old_img_w = old_img_h = imgsz\n",
    "    old_img_b = 1\n",
    "\n",
    "    t0 = time.time()\n",
    "    for path, img, im0s, vid_cap in dataset:\n",
    "        img = torch.from_numpy(img).to(device)\n",
    "        img = img.half() if half else img.float()  # uint8 to fp16/32\n",
    "        img /= 255.0  # 0 - 255 to 0.0 - 1.0\n",
    "        if img.ndimension() == 3:\n",
    "            img = img.unsqueeze(0)\n",
    "\n",
    "        # Warmup\n",
    "        if device.type != 'cpu' and (old_img_b != img.shape[0] or old_img_h != img.shape[2] or old_img_w != img.shape[3]):\n",
    "            old_img_b = img.shape[0]\n",
    "            old_img_h = img.shape[2]\n",
    "            old_img_w = img.shape[3]\n",
    "            for i in range(3):\n",
    "                model(img, augment=augment)[0]\n",
    "\n",
    "        # Inference\n",
    "        t1 = time_synchronized()\n",
    "        with torch.no_grad():   # Calculating gradients would cause a GPU memory leak\n",
    "            pred = model(img, augment=augment)[0]\n",
    "        t2 = time_synchronized()\n",
    "\n",
    "        # Apply NMS\n",
    "        pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=agnostic_nms)\n",
    "        t3 = time_synchronized()\n",
    "\n",
    "        # Process detections\n",
    "        for i, det in enumerate(pred):  # detections per image\n",
    "            p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)\n",
    "\n",
    "            p = Path(p)  # to Path\n",
    "            save_path = p.name # img.jpg\n",
    "            gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh\n",
    "            if len(det):\n",
    "                # Rescale boxes from img_size to im0 size\n",
    "                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()\n",
    "\n",
    "                # Print results\n",
    "                for c in det[:, -1].unique():\n",
    "                    n = (det[:, -1] == c).sum()  # detections per class\n",
    "                    s += f\"{n} {names[int(c)]}{'s' * (n > 1)}, \"  # add to string\n",
    "\n",
    "                # Write results\n",
    "                for *xyxy, conf, cls in reversed(det):\n",
    "                    if save_img or view_img:  # Add bbox to image\n",
    "                        label = f'{names[int(cls)]} {conf:.2f}'\n",
    "                        plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=1)\n",
    "\n",
    "            # Print time (inference + NMS)\n",
    "            print(f'{s}Done. ({(1E3 * (t2 - t1)):.1f}ms) Inference, ({(1E3 * (t3 - t2)):.1f}ms) NMS')\n",
    "\n",
    "            # Stream result\n",
    "            if view_img:\n",
    "                cv2.imshow(str(p), im0)\n",
    "                cv2.waitKey(0)\n",
    "                cv2.destroyAllWindows()\n",
    "\n",
    "            # Save result (image with detections)\n",
    "            if save_img:\n",
    "                if dataset.mode == 'image':\n",
    "                    cv2.imwrite(save_path, cv2.cvtColor(im0, cv2.COLOR_BGR2RGB))\n",
    "                    print(f' The image with the result is saved in: {save_path}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "ee18e2fb-a1c4-4972-9342-f3d6269d315f",
   "metadata": {},
   "outputs": [],
   "source": [
    "source = 'inference/images/bus.jpg' # file\n",
    "img_size = 640  # inference size (pixels)\n",
    "\n",
    "conf_thres = 0.25 # object confidence threshold\n",
    "iou_thres = 0.45 # IOU threshold for NMS\n",
    "\n",
    "view_img = True # display result\n",
    "save_img = False # save result\n",
    "\n",
    "augment = False # augmented inference\n",
    "classes = None # filter by class: 0, or 0 2 3\n",
    "agnostic_nms = False # class-agnostic NMS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "9251f326-5d5e-4310-8ef7-b3d2306102a4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "4 persons, 1 bus, 1 tie, Done. (13.0ms) Inference, (1.0ms) NMS\n"
     ]
    }
   ],
   "source": [
    "detect(source, view_img, save_img)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "646d3977-b5a0-4992-9382-d58d27373016",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-1-12",
   "language": "python",
   "name": "pytorch-1-12"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}