File size: 2,947 Bytes
a33e03b 360a2c0 862884c 360a2c0 862884c 0d49177 360a2c0 862884c 360a2c0 819890a 360a2c0 b6b57c7 360a2c0 b6b57c7 819890a 0d49177 862884c 360a2c0 b6b57c7 862884c 360a2c0 862884c 819890a b6b57c7 360a2c0 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 |
{
"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 (\n",
" AugmentationComposer,\n",
" Config,\n",
" PostProccess,\n",
" create_converter,\n",
" create_model,\n",
" custom_logger,\n",
" draw_bboxes,\n",
")"
]
},
{
"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 = create_converter(cfg.model.name, 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",
""
]
}
],
"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
}
|