File size: 3,233 Bytes
a33e03b 360a2c0 862884c 360a2c0 862884c 360a2c0 862884c 360a2c0 819890a 360a2c0 b6b57c7 360a2c0 b6b57c7 819890a 862884c 360a2c0 b6b57c7 862884c 360a2c0 862884c 819890a b6b57c7 360a2c0 862884c a33e03b 360a2c0 a33e03b 360a2c0 a33e03b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import sys\n",
"from pathlib import Path\n",
"\n",
"import torch\n",
"from hydra import compose, initialize\n",
"from PIL import Image \n",
"\n",
"project_root = Path().resolve().parent\n",
"sys.path.append(str(project_root))\n",
"\n",
"from yolo import AugmentationComposer, Config, create_model, custom_logger, draw_bboxes, Vec2Box, PostProccess\n",
"from yolo.utils.bounding_box_utils import Anc2Box"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"CONFIG_PATH = \"../yolo/config\"\n",
"CONFIG_NAME = \"config\"\n",
"MODEL = \"v7-base\"\n",
"\n",
"DEVICE = 'cuda:0'\n",
"CLASS_NUM = 80\n",
"IMAGE_PATH = '../demo/images/inference/image.png'\n",
"\n",
"custom_logger()\n",
"device = torch.device(DEVICE)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with initialize(config_path=CONFIG_PATH, version_base=None, job_name=\"notebook_job\"):\n",
" cfg: Config = compose(config_name=CONFIG_NAME, overrides=[\"task=inference\", f\"task.data.source={IMAGE_PATH}\", f\"model={MODEL}\"])\n",
" model = create_model(cfg.model, class_num=CLASS_NUM).to(device)\n",
" transform = AugmentationComposer([], cfg.image_size)\n",
" converter = Anc2Box(model, cfg.model.anchor, cfg.image_size, device)\n",
" # converter = Vec2Box(model, cfg.model.anchor, cfg.image_size, device)\n",
" post_proccess = PostProccess(converter, cfg.task.nms)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"pil_image = Image.open(IMAGE_PATH)\n",
"image, bbox, rev_tensor = transform(pil_image)\n",
"image = image.to(device)[None]\n",
"rev_tensor = rev_tensor.to(device)[None]"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"with torch.no_grad():\n",
" predict = model(image)\n",
" pred_bbox = post_proccess(predict, rev_tensor)\n",
"\n",
"draw_bboxes(pil_image, pred_bbox, idx2label=cfg.class_list)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Sample Output:\n",
"\n",
""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%load_ext autoreload\n",
"%autoreload 2"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "yolomit",
"language": "python",
"name": "python3"
},
"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.10.14"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
|