henry000 commited on
Commit
360a2c0
·
1 Parent(s): 40afe67

✨ [Add] an example ipynb! demo how to inference

Browse files
.gitignore CHANGED
@@ -145,4 +145,4 @@ runs
145
  node_modules/
146
 
147
  # Not ignore image for demo
148
- !demo/images/inference/*
 
145
  node_modules/
146
 
147
  # Not ignore image for demo
148
+ !demo/images/*/*.png
demo/images/output/visualize.png ADDED
examples/notebook_inference.ipynb CHANGED
@@ -5,7 +5,76 @@
5
  "execution_count": null,
6
  "metadata": {},
7
  "outputs": [],
8
- "source": []
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
9
  },
10
  {
11
  "cell_type": "code",
@@ -16,8 +85,22 @@
16
  }
17
  ],
18
  "metadata": {
 
 
 
 
 
19
  "language_info": {
20
- "name": "python"
 
 
 
 
 
 
 
 
 
21
  }
22
  },
23
  "nbformat": 4,
 
5
  "execution_count": null,
6
  "metadata": {},
7
  "outputs": [],
8
+ "source": [
9
+ "import torch\n",
10
+ "from hydra import compose, initialize\n",
11
+ "from PIL import Image \n",
12
+ "\n",
13
+ "from yolo import AugmentationComposer, bbox_nms, Config, create_model, custom_logger, draw_bboxes"
14
+ ]
15
+ },
16
+ {
17
+ "cell_type": "code",
18
+ "execution_count": null,
19
+ "metadata": {},
20
+ "outputs": [],
21
+ "source": [
22
+ "CONFIG_PATH = \"../yolo/config\"\n",
23
+ "CONFIG_NAME = \"config\"\n",
24
+ "\n",
25
+ "DEVICE = 'cuda:0'\n",
26
+ "# If weight not exist there, it will try to download it\n",
27
+ "\n",
28
+ "WEIGHT_PATH = '../weights/v9-cnw.pt' \n",
29
+ "IMAGE_PATH = '../demo/images/inference/image.png'\n",
30
+ "\n",
31
+ "custom_logger()\n",
32
+ "device = torch.device(DEVICE)\n",
33
+ "image = Image.open(IMAGE_PATH)"
34
+ ]
35
+ },
36
+ {
37
+ "cell_type": "code",
38
+ "execution_count": null,
39
+ "metadata": {},
40
+ "outputs": [],
41
+ "source": [
42
+ "with initialize(config_path=CONFIG_PATH, version_base=None, job_name=\"notebook_job\"):\n",
43
+ " cfg: Config = compose(config_name=CONFIG_NAME, overrides=[\"task=inference\", f\"task.data.source={IMAGE_PATH}\"])\n",
44
+ " model = create_model(cfg.model, WEIGHT_PATH).to(device)\n",
45
+ " transform = AugmentationComposer([], cfg.image_size)"
46
+ ]
47
+ },
48
+ {
49
+ "cell_type": "code",
50
+ "execution_count": null,
51
+ "metadata": {},
52
+ "outputs": [],
53
+ "source": [
54
+ "image, bbox = transform(image, torch.zeros(0, 5))\n",
55
+ "image = image.to(device)[None]"
56
+ ]
57
+ },
58
+ {
59
+ "cell_type": "code",
60
+ "execution_count": null,
61
+ "metadata": {},
62
+ "outputs": [],
63
+ "source": [
64
+ "with torch.no_grad():\n",
65
+ " predict = model(image)\n",
66
+ "predict_box = bbox_nms(predict[1][0], cfg.task.nms)\n",
67
+ "draw_bboxes(image, predict_box, save_path='../demo/images/output/', idx2label=cfg.class_list)"
68
+ ]
69
+ },
70
+ {
71
+ "cell_type": "markdown",
72
+ "metadata": {},
73
+ "source": [
74
+ "Sample Output:\n",
75
+ "\n",
76
+ "![image](../demo/images/output/visualize.png)"
77
+ ]
78
  },
79
  {
80
  "cell_type": "code",
 
85
  }
86
  ],
87
  "metadata": {
88
+ "kernelspec": {
89
+ "display_name": "yolomit",
90
+ "language": "python",
91
+ "name": "python3"
92
+ },
93
  "language_info": {
94
+ "codemirror_mode": {
95
+ "name": "ipython",
96
+ "version": 3
97
+ },
98
+ "file_extension": ".py",
99
+ "mimetype": "text/x-python",
100
+ "name": "python",
101
+ "nbconvert_exporter": "python",
102
+ "pygments_lexer": "ipython3",
103
+ "version": "3.10.14"
104
  }
105
  },
106
  "nbformat": 4,