dental-detector 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- dental_detector-0.1.0/.github/workflows/tests.yml +59 -0
- dental_detector-0.1.0/LICENSE +21 -0
- dental_detector-0.1.0/PKG-INFO +196 -0
- dental_detector-0.1.0/README.md +159 -0
- dental_detector-0.1.0/dental_detector/__init__.py +19 -0
- dental_detector-0.1.0/dental_detector/cli.py +73 -0
- dental_detector-0.1.0/dental_detector/detector.py +177 -0
- dental_detector-0.1.0/dental_detector/models.py +111 -0
- dental_detector-0.1.0/dental_detector/utils.py +62 -0
- dental_detector-0.1.0/dental_detector/visualization.py +61 -0
- dental_detector-0.1.0/pyproject.toml +77 -0
- dental_detector-0.1.0/tests/__init__.py +0 -0
- dental_detector-0.1.0/tests/test_detector.py +344 -0
|
@@ -0,0 +1,59 @@
|
|
|
1
|
+
name: Tests
|
|
2
|
+
|
|
3
|
+
on:
|
|
4
|
+
push:
|
|
5
|
+
branches: [main, master]
|
|
6
|
+
pull_request:
|
|
7
|
+
branches: [main, master]
|
|
8
|
+
|
|
9
|
+
jobs:
|
|
10
|
+
test:
|
|
11
|
+
runs-on: ubuntu-latest
|
|
12
|
+
strategy:
|
|
13
|
+
matrix:
|
|
14
|
+
python-version: ["3.9", "3.10", "3.11", "3.12"]
|
|
15
|
+
|
|
16
|
+
steps:
|
|
17
|
+
- uses: actions/checkout@v4
|
|
18
|
+
|
|
19
|
+
- name: Set up Python ${{ matrix.python-version }}
|
|
20
|
+
uses: actions/setup-python@v5
|
|
21
|
+
with:
|
|
22
|
+
python-version: ${{ matrix.python-version }}
|
|
23
|
+
|
|
24
|
+
- name: Install dependencies
|
|
25
|
+
run: |
|
|
26
|
+
python -m pip install --upgrade pip
|
|
27
|
+
pip install -e ".[dev]"
|
|
28
|
+
|
|
29
|
+
- name: Lint with ruff
|
|
30
|
+
run: ruff check dental_detector tests
|
|
31
|
+
|
|
32
|
+
- name: Run tests
|
|
33
|
+
run: pytest --cov=dental_detector --cov-report=term-missing
|
|
34
|
+
|
|
35
|
+
publish:
|
|
36
|
+
name: Publish to PyPI
|
|
37
|
+
needs: test
|
|
38
|
+
runs-on: ubuntu-latest
|
|
39
|
+
if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/v')
|
|
40
|
+
|
|
41
|
+
permissions:
|
|
42
|
+
id-token: write
|
|
43
|
+
|
|
44
|
+
steps:
|
|
45
|
+
- uses: actions/checkout@v4
|
|
46
|
+
|
|
47
|
+
- name: Set up Python
|
|
48
|
+
uses: actions/setup-python@v5
|
|
49
|
+
with:
|
|
50
|
+
python-version: "3.11"
|
|
51
|
+
|
|
52
|
+
- name: Install build tools
|
|
53
|
+
run: pip install hatchling build
|
|
54
|
+
|
|
55
|
+
- name: Build package
|
|
56
|
+
run: python -m build
|
|
57
|
+
|
|
58
|
+
- name: Publish to PyPI
|
|
59
|
+
uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2025 Sandeep Vissa
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,196 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: dental-detector
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Tooth detection in maxilla and mandible dental images using YOLOv8
|
|
5
|
+
Project-URL: Homepage, https://github.com/sandeepvissa/dental-detector
|
|
6
|
+
Project-URL: Repository, https://github.com/sandeepvissa/dental-detector
|
|
7
|
+
Project-URL: Bug Tracker, https://github.com/sandeepvissa/dental-detector/issues
|
|
8
|
+
Author-email: Sandeep Vissa <sandeepvissa2002@gmail.com>
|
|
9
|
+
License: MIT
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Keywords: computer-vision,dental,dental-imaging,mandible,maxilla,object-detection,teeth,tooth-detection,yolo
|
|
12
|
+
Classifier: Development Status :: 3 - Alpha
|
|
13
|
+
Classifier: Intended Audience :: Developers
|
|
14
|
+
Classifier: Intended Audience :: Healthcare Industry
|
|
15
|
+
Classifier: Intended Audience :: Science/Research
|
|
16
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
17
|
+
Classifier: Operating System :: OS Independent
|
|
18
|
+
Classifier: Programming Language :: Python :: 3
|
|
19
|
+
Classifier: Programming Language :: Python :: 3.8
|
|
20
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
21
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
22
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
23
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
24
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
25
|
+
Classifier: Topic :: Scientific/Engineering :: Image Recognition
|
|
26
|
+
Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
|
|
27
|
+
Requires-Python: >=3.8
|
|
28
|
+
Requires-Dist: numpy>=1.21.0
|
|
29
|
+
Requires-Dist: opencv-python>=4.5.0
|
|
30
|
+
Requires-Dist: pillow>=9.0.0
|
|
31
|
+
Requires-Dist: ultralytics>=8.0.0
|
|
32
|
+
Provides-Extra: dev
|
|
33
|
+
Requires-Dist: pytest-cov>=4.0; extra == 'dev'
|
|
34
|
+
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
35
|
+
Requires-Dist: ruff>=0.1.0; extra == 'dev'
|
|
36
|
+
Description-Content-Type: text/markdown
|
|
37
|
+
|
|
38
|
+
# dental-detector
|
|
39
|
+
|
|
40
|
+
[](https://pypi.org/project/dental-detector/)
|
|
41
|
+
[](https://pypi.org/project/dental-detector/)
|
|
42
|
+
[](LICENSE)
|
|
43
|
+
[](https://github.com/sandeepvissa/dental-detector/actions/workflows/tests.yml)
|
|
44
|
+
|
|
45
|
+
**Tooth detection in maxilla and mandible dental images using YOLOv8.**
|
|
46
|
+
|
|
47
|
+
`dental-detector` is a lightweight, open-source Python package that wraps a YOLOv8 model trained to locate and classify teeth in dental radiographs and photographs. It detects 7 tooth types across both arches and ships with a clean Python API and a CLI.
|
|
48
|
+
|
|
49
|
+
## Detected classes
|
|
50
|
+
|
|
51
|
+
| Class ID | Tooth |
|
|
52
|
+
|----------|-------|
|
|
53
|
+
| 0 | 1st Molar |
|
|
54
|
+
| 1 | 1st Premolar |
|
|
55
|
+
| 2 | 2nd Molar |
|
|
56
|
+
| 3 | 2nd Premolar |
|
|
57
|
+
| 4 | Canine |
|
|
58
|
+
| 5 | Central Incisor |
|
|
59
|
+
| 6 | Lateral Incisor |
|
|
60
|
+
|
|
61
|
+
## Installation
|
|
62
|
+
|
|
63
|
+
```bash
|
|
64
|
+
pip install dental-detector
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
> **Note:** You must supply your own trained YOLOv8 weights (`.pt` file).
|
|
68
|
+
> The package does not bundle a model; it provides the inference wrapper.
|
|
69
|
+
|
|
70
|
+
## Quick start
|
|
71
|
+
|
|
72
|
+
```python
|
|
73
|
+
from dental_detector import DentalDetector
|
|
74
|
+
|
|
75
|
+
detector = DentalDetector("best.pt")
|
|
76
|
+
|
|
77
|
+
# --- Detect ---
|
|
78
|
+
result = detector.detect("xray_maxilla.jpg")
|
|
79
|
+
print(f"Found {len(result)} teeth")
|
|
80
|
+
for tooth in result:
|
|
81
|
+
print(tooth.label, tooth.confidence, tooth.bbox)
|
|
82
|
+
|
|
83
|
+
# --- Annotate and save ---
|
|
84
|
+
annotated = detector.annotate("xray_maxilla.jpg") # PIL Image
|
|
85
|
+
annotated.save("annotated.jpg")
|
|
86
|
+
|
|
87
|
+
# or in one step:
|
|
88
|
+
detector.annotate_and_save("xray_maxilla.jpg", "annotated.jpg")
|
|
89
|
+
|
|
90
|
+
# --- Export as JSON ---
|
|
91
|
+
json_str = detector.detect_to_json("xray_mandible.jpg")
|
|
92
|
+
print(json_str)
|
|
93
|
+
```
|
|
94
|
+
|
|
95
|
+
### Working with PIL / numpy directly
|
|
96
|
+
|
|
97
|
+
```python
|
|
98
|
+
from PIL import Image
|
|
99
|
+
import numpy as np
|
|
100
|
+
|
|
101
|
+
pil_img = Image.open("xray.jpg")
|
|
102
|
+
result = detector.detect(pil_img)
|
|
103
|
+
|
|
104
|
+
arr = np.array(pil_img) # RGB numpy array also works
|
|
105
|
+
result = detector.detect(arr)
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
### Confidence threshold
|
|
109
|
+
|
|
110
|
+
```python
|
|
111
|
+
# Set at construction time
|
|
112
|
+
detector = DentalDetector("best.pt", conf_threshold=0.7)
|
|
113
|
+
|
|
114
|
+
# Or override per call
|
|
115
|
+
result = detector.detect("xray.jpg", conf_threshold=0.5)
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
### Filtering results
|
|
119
|
+
|
|
120
|
+
```python
|
|
121
|
+
# Keep only molars
|
|
122
|
+
molars = result.filter_by_label("1st Molar")
|
|
123
|
+
|
|
124
|
+
# Keep high-confidence detections
|
|
125
|
+
high_conf = result.filter_by_confidence(0.85)
|
|
126
|
+
|
|
127
|
+
# Unique tooth types present in the image
|
|
128
|
+
print(result.unique_labels)
|
|
129
|
+
```
|
|
130
|
+
|
|
131
|
+
## Command-line interface
|
|
132
|
+
|
|
133
|
+
```bash
|
|
134
|
+
# Print detected teeth
|
|
135
|
+
dental-detector xray.jpg --model best.pt
|
|
136
|
+
|
|
137
|
+
# Save annotated image
|
|
138
|
+
dental-detector xray.jpg --model best.pt --output annotated.jpg
|
|
139
|
+
|
|
140
|
+
# Output as JSON
|
|
141
|
+
dental-detector xray.jpg --model best.pt --json
|
|
142
|
+
|
|
143
|
+
# Set confidence threshold
|
|
144
|
+
dental-detector xray.jpg --model best.pt --conf 0.7
|
|
145
|
+
```
|
|
146
|
+
|
|
147
|
+
## API reference
|
|
148
|
+
|
|
149
|
+
### `DentalDetector`
|
|
150
|
+
|
|
151
|
+
| Method | Returns | Description |
|
|
152
|
+
|--------|---------|-------------|
|
|
153
|
+
| `detect(image, conf_threshold=None)` | `DetectionResult` | Run detection |
|
|
154
|
+
| `annotate(image, ...)` | `PIL.Image` | Detect + draw boxes |
|
|
155
|
+
| `annotate_and_save(image, path, ...)` | `Path` | Detect, draw, save |
|
|
156
|
+
| `detect_to_json(image, ...)` | `str` | Detections as JSON |
|
|
157
|
+
|
|
158
|
+
### `DetectionResult`
|
|
159
|
+
|
|
160
|
+
| Attribute / Method | Description |
|
|
161
|
+
|--------------------|-------------|
|
|
162
|
+
| `detections` | `list[Detection]` |
|
|
163
|
+
| `labels` | All label strings |
|
|
164
|
+
| `unique_labels` | Sorted unique labels |
|
|
165
|
+
| `filter_by_label(label)` | Returns filtered `DetectionResult` |
|
|
166
|
+
| `filter_by_confidence(threshold)` | Returns filtered `DetectionResult` |
|
|
167
|
+
| `to_dict()` | Serializable dict |
|
|
168
|
+
|
|
169
|
+
### `Detection`
|
|
170
|
+
|
|
171
|
+
| Attribute | Type | Description |
|
|
172
|
+
|-----------|------|-------------|
|
|
173
|
+
| `label` | `str` | Tooth class name |
|
|
174
|
+
| `class_id` | `int` | Class index (0–6) |
|
|
175
|
+
| `confidence` | `float` | Model confidence (0–1) |
|
|
176
|
+
| `bbox` | `tuple[float, float, float, float]` | `(x1, y1, x2, y2)` pixels |
|
|
177
|
+
| `center` | `tuple[float, float]` | Bounding box center |
|
|
178
|
+
| `width` / `height` / `area` | `float` | Bounding box dimensions |
|
|
179
|
+
|
|
180
|
+
## Development
|
|
181
|
+
|
|
182
|
+
```bash
|
|
183
|
+
git clone https://github.com/sandeepvissa/dental-detector
|
|
184
|
+
cd dental-detector
|
|
185
|
+
pip install -e ".[dev]"
|
|
186
|
+
|
|
187
|
+
# Run tests
|
|
188
|
+
pytest
|
|
189
|
+
|
|
190
|
+
# Lint
|
|
191
|
+
ruff check dental_detector tests
|
|
192
|
+
```
|
|
193
|
+
|
|
194
|
+
## License
|
|
195
|
+
|
|
196
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,159 @@
|
|
|
1
|
+
# dental-detector
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/dental-detector/)
|
|
4
|
+
[](https://pypi.org/project/dental-detector/)
|
|
5
|
+
[](LICENSE)
|
|
6
|
+
[](https://github.com/sandeepvissa/dental-detector/actions/workflows/tests.yml)
|
|
7
|
+
|
|
8
|
+
**Tooth detection in maxilla and mandible dental images using YOLOv8.**
|
|
9
|
+
|
|
10
|
+
`dental-detector` is a lightweight, open-source Python package that wraps a YOLOv8 model trained to locate and classify teeth in dental radiographs and photographs. It detects 7 tooth types across both arches and ships with a clean Python API and a CLI.
|
|
11
|
+
|
|
12
|
+
## Detected classes
|
|
13
|
+
|
|
14
|
+
| Class ID | Tooth |
|
|
15
|
+
|----------|-------|
|
|
16
|
+
| 0 | 1st Molar |
|
|
17
|
+
| 1 | 1st Premolar |
|
|
18
|
+
| 2 | 2nd Molar |
|
|
19
|
+
| 3 | 2nd Premolar |
|
|
20
|
+
| 4 | Canine |
|
|
21
|
+
| 5 | Central Incisor |
|
|
22
|
+
| 6 | Lateral Incisor |
|
|
23
|
+
|
|
24
|
+
## Installation
|
|
25
|
+
|
|
26
|
+
```bash
|
|
27
|
+
pip install dental-detector
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
> **Note:** You must supply your own trained YOLOv8 weights (`.pt` file).
|
|
31
|
+
> The package does not bundle a model; it provides the inference wrapper.
|
|
32
|
+
|
|
33
|
+
## Quick start
|
|
34
|
+
|
|
35
|
+
```python
|
|
36
|
+
from dental_detector import DentalDetector
|
|
37
|
+
|
|
38
|
+
detector = DentalDetector("best.pt")
|
|
39
|
+
|
|
40
|
+
# --- Detect ---
|
|
41
|
+
result = detector.detect("xray_maxilla.jpg")
|
|
42
|
+
print(f"Found {len(result)} teeth")
|
|
43
|
+
for tooth in result:
|
|
44
|
+
print(tooth.label, tooth.confidence, tooth.bbox)
|
|
45
|
+
|
|
46
|
+
# --- Annotate and save ---
|
|
47
|
+
annotated = detector.annotate("xray_maxilla.jpg") # PIL Image
|
|
48
|
+
annotated.save("annotated.jpg")
|
|
49
|
+
|
|
50
|
+
# or in one step:
|
|
51
|
+
detector.annotate_and_save("xray_maxilla.jpg", "annotated.jpg")
|
|
52
|
+
|
|
53
|
+
# --- Export as JSON ---
|
|
54
|
+
json_str = detector.detect_to_json("xray_mandible.jpg")
|
|
55
|
+
print(json_str)
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
### Working with PIL / numpy directly
|
|
59
|
+
|
|
60
|
+
```python
|
|
61
|
+
from PIL import Image
|
|
62
|
+
import numpy as np
|
|
63
|
+
|
|
64
|
+
pil_img = Image.open("xray.jpg")
|
|
65
|
+
result = detector.detect(pil_img)
|
|
66
|
+
|
|
67
|
+
arr = np.array(pil_img) # RGB numpy array also works
|
|
68
|
+
result = detector.detect(arr)
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
### Confidence threshold
|
|
72
|
+
|
|
73
|
+
```python
|
|
74
|
+
# Set at construction time
|
|
75
|
+
detector = DentalDetector("best.pt", conf_threshold=0.7)
|
|
76
|
+
|
|
77
|
+
# Or override per call
|
|
78
|
+
result = detector.detect("xray.jpg", conf_threshold=0.5)
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
### Filtering results
|
|
82
|
+
|
|
83
|
+
```python
|
|
84
|
+
# Keep only molars
|
|
85
|
+
molars = result.filter_by_label("1st Molar")
|
|
86
|
+
|
|
87
|
+
# Keep high-confidence detections
|
|
88
|
+
high_conf = result.filter_by_confidence(0.85)
|
|
89
|
+
|
|
90
|
+
# Unique tooth types present in the image
|
|
91
|
+
print(result.unique_labels)
|
|
92
|
+
```
|
|
93
|
+
|
|
94
|
+
## Command-line interface
|
|
95
|
+
|
|
96
|
+
```bash
|
|
97
|
+
# Print detected teeth
|
|
98
|
+
dental-detector xray.jpg --model best.pt
|
|
99
|
+
|
|
100
|
+
# Save annotated image
|
|
101
|
+
dental-detector xray.jpg --model best.pt --output annotated.jpg
|
|
102
|
+
|
|
103
|
+
# Output as JSON
|
|
104
|
+
dental-detector xray.jpg --model best.pt --json
|
|
105
|
+
|
|
106
|
+
# Set confidence threshold
|
|
107
|
+
dental-detector xray.jpg --model best.pt --conf 0.7
|
|
108
|
+
```
|
|
109
|
+
|
|
110
|
+
## API reference
|
|
111
|
+
|
|
112
|
+
### `DentalDetector`
|
|
113
|
+
|
|
114
|
+
| Method | Returns | Description |
|
|
115
|
+
|--------|---------|-------------|
|
|
116
|
+
| `detect(image, conf_threshold=None)` | `DetectionResult` | Run detection |
|
|
117
|
+
| `annotate(image, ...)` | `PIL.Image` | Detect + draw boxes |
|
|
118
|
+
| `annotate_and_save(image, path, ...)` | `Path` | Detect, draw, save |
|
|
119
|
+
| `detect_to_json(image, ...)` | `str` | Detections as JSON |
|
|
120
|
+
|
|
121
|
+
### `DetectionResult`
|
|
122
|
+
|
|
123
|
+
| Attribute / Method | Description |
|
|
124
|
+
|--------------------|-------------|
|
|
125
|
+
| `detections` | `list[Detection]` |
|
|
126
|
+
| `labels` | All label strings |
|
|
127
|
+
| `unique_labels` | Sorted unique labels |
|
|
128
|
+
| `filter_by_label(label)` | Returns filtered `DetectionResult` |
|
|
129
|
+
| `filter_by_confidence(threshold)` | Returns filtered `DetectionResult` |
|
|
130
|
+
| `to_dict()` | Serializable dict |
|
|
131
|
+
|
|
132
|
+
### `Detection`
|
|
133
|
+
|
|
134
|
+
| Attribute | Type | Description |
|
|
135
|
+
|-----------|------|-------------|
|
|
136
|
+
| `label` | `str` | Tooth class name |
|
|
137
|
+
| `class_id` | `int` | Class index (0–6) |
|
|
138
|
+
| `confidence` | `float` | Model confidence (0–1) |
|
|
139
|
+
| `bbox` | `tuple[float, float, float, float]` | `(x1, y1, x2, y2)` pixels |
|
|
140
|
+
| `center` | `tuple[float, float]` | Bounding box center |
|
|
141
|
+
| `width` / `height` / `area` | `float` | Bounding box dimensions |
|
|
142
|
+
|
|
143
|
+
## Development
|
|
144
|
+
|
|
145
|
+
```bash
|
|
146
|
+
git clone https://github.com/sandeepvissa/dental-detector
|
|
147
|
+
cd dental-detector
|
|
148
|
+
pip install -e ".[dev]"
|
|
149
|
+
|
|
150
|
+
# Run tests
|
|
151
|
+
pytest
|
|
152
|
+
|
|
153
|
+
# Lint
|
|
154
|
+
ruff check dental_detector tests
|
|
155
|
+
```
|
|
156
|
+
|
|
157
|
+
## License
|
|
158
|
+
|
|
159
|
+
MIT — see [LICENSE](LICENSE).
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
"""dental_detector — tooth detection in maxilla and mandible X-ray images."""
|
|
2
|
+
|
|
3
|
+
from .detector import DentalDetector
|
|
4
|
+
from .models import Detection, DetectionResult, TOOTH_CLASSES, LABEL_COLORS
|
|
5
|
+
from .utils import load_image, to_pil, letterbox
|
|
6
|
+
from .visualization import annotate
|
|
7
|
+
|
|
8
|
+
__version__ = "0.1.0"
|
|
9
|
+
__all__ = [
|
|
10
|
+
"DentalDetector",
|
|
11
|
+
"Detection",
|
|
12
|
+
"DetectionResult",
|
|
13
|
+
"TOOTH_CLASSES",
|
|
14
|
+
"LABEL_COLORS",
|
|
15
|
+
"load_image",
|
|
16
|
+
"to_pil",
|
|
17
|
+
"letterbox",
|
|
18
|
+
"annotate",
|
|
19
|
+
]
|
|
@@ -0,0 +1,73 @@
|
|
|
1
|
+
"""Command-line interface for dental_detector."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import argparse
|
|
6
|
+
import json
|
|
7
|
+
import sys
|
|
8
|
+
from pathlib import Path
|
|
9
|
+
|
|
10
|
+
from dental_detector import DentalDetector
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def _build_parser() -> argparse.ArgumentParser:
|
|
14
|
+
p = argparse.ArgumentParser(
|
|
15
|
+
prog="dental-detector",
|
|
16
|
+
description="Detect teeth in dental images (maxilla / mandible) using YOLOv8.",
|
|
17
|
+
)
|
|
18
|
+
p.add_argument("image", help="Path to the input image.")
|
|
19
|
+
p.add_argument("--model", "-m", required=True, help="Path to YOLOv8 .pt weights file.")
|
|
20
|
+
p.add_argument(
|
|
21
|
+
"--conf", "-c", type=float, default=0.6,
|
|
22
|
+
help="Confidence threshold (default: 0.6).",
|
|
23
|
+
)
|
|
24
|
+
p.add_argument(
|
|
25
|
+
"--output", "-o", default=None,
|
|
26
|
+
help="Save annotated image to this path. Skipped if not provided.",
|
|
27
|
+
)
|
|
28
|
+
p.add_argument(
|
|
29
|
+
"--json", action="store_true",
|
|
30
|
+
help="Print detections as JSON instead of plain text.",
|
|
31
|
+
)
|
|
32
|
+
p.add_argument(
|
|
33
|
+
"--no-confidence", action="store_true",
|
|
34
|
+
help="Hide confidence scores in the annotated image.",
|
|
35
|
+
)
|
|
36
|
+
return p
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def main(argv=None) -> int:
|
|
40
|
+
args = _build_parser().parse_args(argv)
|
|
41
|
+
|
|
42
|
+
try:
|
|
43
|
+
detector = DentalDetector(args.model, conf_threshold=args.conf)
|
|
44
|
+
result = detector.detect(args.image)
|
|
45
|
+
except (FileNotFoundError, ValueError, RuntimeError) as exc:
|
|
46
|
+
print(f"Error: {exc}", file=sys.stderr)
|
|
47
|
+
return 1
|
|
48
|
+
|
|
49
|
+
if args.json:
|
|
50
|
+
print(json.dumps(result.to_dict(), indent=2))
|
|
51
|
+
else:
|
|
52
|
+
print(f"Detected {len(result)} tooth/teeth in '{Path(args.image).name}':")
|
|
53
|
+
for det in result:
|
|
54
|
+
x1, y1, x2, y2 = det.bbox
|
|
55
|
+
print(
|
|
56
|
+
f" [{det.class_id}] {det.label:<18} conf={det.confidence:.3f} "
|
|
57
|
+
f"bbox=({x1:.0f},{y1:.0f},{x2:.0f},{y2:.0f})"
|
|
58
|
+
)
|
|
59
|
+
|
|
60
|
+
if args.output:
|
|
61
|
+
out_path = detector.annotate_and_save(
|
|
62
|
+
args.image,
|
|
63
|
+
args.output,
|
|
64
|
+
conf_threshold=args.conf,
|
|
65
|
+
show_confidence=not args.no_confidence,
|
|
66
|
+
)
|
|
67
|
+
print(f"Annotated image saved to: {out_path}")
|
|
68
|
+
|
|
69
|
+
return 0
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
if __name__ == "__main__":
|
|
73
|
+
sys.exit(main())
|
|
@@ -0,0 +1,177 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
from typing import Optional, Union
|
|
6
|
+
|
|
7
|
+
import cv2
|
|
8
|
+
import numpy as np
|
|
9
|
+
from PIL import Image
|
|
10
|
+
|
|
11
|
+
try:
|
|
12
|
+
from ultralytics import YOLO
|
|
13
|
+
except ImportError:
|
|
14
|
+
YOLO = None # type: ignore[assignment,misc]
|
|
15
|
+
|
|
16
|
+
from .models import Detection, DetectionResult, TOOTH_CLASSES
|
|
17
|
+
from .utils import ImageInput, load_image, to_pil
|
|
18
|
+
from .visualization import annotate as _annotate
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class DentalDetector:
|
|
22
|
+
"""Detect teeth in dental X-ray or photograph images using a YOLOv8 model.
|
|
23
|
+
|
|
24
|
+
Supported jaw views: maxilla (upper) and mandible (lower).
|
|
25
|
+
|
|
26
|
+
Args:
|
|
27
|
+
model_path: Path to a YOLOv8 ``.pt`` weights file.
|
|
28
|
+
conf_threshold: Default minimum confidence score (0–1). Detections
|
|
29
|
+
below this value are discarded. Can be overridden per call.
|
|
30
|
+
device: Inference device string understood by PyTorch, e.g. ``"cpu"``,
|
|
31
|
+
``"cuda"``, ``"cuda:0"``. ``None`` auto-selects GPU when available.
|
|
32
|
+
|
|
33
|
+
Example::
|
|
34
|
+
|
|
35
|
+
from dental_detector import DentalDetector
|
|
36
|
+
|
|
37
|
+
detector = DentalDetector("best.pt")
|
|
38
|
+
result = detector.detect("xray.jpg")
|
|
39
|
+
for tooth in result:
|
|
40
|
+
print(tooth.label, tooth.confidence, tooth.bbox)
|
|
41
|
+
|
|
42
|
+
annotated = detector.annotate("xray.jpg")
|
|
43
|
+
annotated.save("output.jpg")
|
|
44
|
+
"""
|
|
45
|
+
|
|
46
|
+
def __init__(
|
|
47
|
+
self,
|
|
48
|
+
model_path: Union[str, Path],
|
|
49
|
+
conf_threshold: float = 0.6,
|
|
50
|
+
device: Optional[str] = None,
|
|
51
|
+
) -> None:
|
|
52
|
+
if YOLO is None:
|
|
53
|
+
raise ImportError("ultralytics is required: pip install ultralytics")
|
|
54
|
+
|
|
55
|
+
model_path = Path(model_path)
|
|
56
|
+
if not model_path.exists():
|
|
57
|
+
raise FileNotFoundError(f"Model weights not found: {model_path}")
|
|
58
|
+
|
|
59
|
+
self._model = YOLO(str(model_path))
|
|
60
|
+
if device is not None:
|
|
61
|
+
self._model.to(device)
|
|
62
|
+
|
|
63
|
+
self.conf_threshold = conf_threshold
|
|
64
|
+
self.model_path = model_path
|
|
65
|
+
|
|
66
|
+
# ------------------------------------------------------------------
|
|
67
|
+
# Detection
|
|
68
|
+
# ------------------------------------------------------------------
|
|
69
|
+
|
|
70
|
+
def detect(
|
|
71
|
+
self,
|
|
72
|
+
image: ImageInput,
|
|
73
|
+
conf_threshold: Optional[float] = None,
|
|
74
|
+
) -> DetectionResult:
|
|
75
|
+
"""Run tooth detection on *image*.
|
|
76
|
+
|
|
77
|
+
Args:
|
|
78
|
+
image: File path, PIL Image, or numpy array (RGB or BGR).
|
|
79
|
+
conf_threshold: Override the instance-level threshold for this call.
|
|
80
|
+
|
|
81
|
+
Returns:
|
|
82
|
+
:class:`DetectionResult` containing all detections above threshold.
|
|
83
|
+
"""
|
|
84
|
+
threshold = conf_threshold if conf_threshold is not None else self.conf_threshold
|
|
85
|
+
bgr = load_image(image)
|
|
86
|
+
h, w = bgr.shape[:2]
|
|
87
|
+
|
|
88
|
+
rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
|
|
89
|
+
raw = self._model(rgb, verbose=False)
|
|
90
|
+
|
|
91
|
+
detections = []
|
|
92
|
+
for r in raw:
|
|
93
|
+
boxes = r.boxes.xyxy.cpu().numpy()
|
|
94
|
+
confs = r.boxes.conf.cpu().numpy()
|
|
95
|
+
classes = r.boxes.cls.cpu().numpy()
|
|
96
|
+
names = r.names
|
|
97
|
+
|
|
98
|
+
for (x1, y1, x2, y2), conf, cls_id in zip(boxes, confs, classes):
|
|
99
|
+
if conf < threshold:
|
|
100
|
+
continue
|
|
101
|
+
cid = int(cls_id)
|
|
102
|
+
label = names.get(cid, TOOTH_CLASSES.get(cid, str(cid)))
|
|
103
|
+
detections.append(
|
|
104
|
+
Detection(
|
|
105
|
+
label=label,
|
|
106
|
+
class_id=cid,
|
|
107
|
+
confidence=float(conf),
|
|
108
|
+
bbox=(float(x1), float(y1), float(x2), float(y2)),
|
|
109
|
+
)
|
|
110
|
+
)
|
|
111
|
+
|
|
112
|
+
return DetectionResult(detections=detections, image_width=w, image_height=h)
|
|
113
|
+
|
|
114
|
+
# ------------------------------------------------------------------
|
|
115
|
+
# Annotation helpers
|
|
116
|
+
# ------------------------------------------------------------------
|
|
117
|
+
|
|
118
|
+
def annotate(
|
|
119
|
+
self,
|
|
120
|
+
image: ImageInput,
|
|
121
|
+
conf_threshold: Optional[float] = None,
|
|
122
|
+
show_confidence: bool = True,
|
|
123
|
+
) -> Image.Image:
|
|
124
|
+
"""Detect and draw bounding boxes, returning a PIL Image.
|
|
125
|
+
|
|
126
|
+
Args:
|
|
127
|
+
image: Source image.
|
|
128
|
+
conf_threshold: Minimum confidence (uses instance default if None).
|
|
129
|
+
show_confidence: Whether to include confidence scores in labels.
|
|
130
|
+
|
|
131
|
+
Returns:
|
|
132
|
+
PIL RGB Image with bounding boxes overlaid.
|
|
133
|
+
"""
|
|
134
|
+
bgr = load_image(image)
|
|
135
|
+
result = self.detect(image, conf_threshold=conf_threshold)
|
|
136
|
+
annotated_bgr = _annotate(bgr, result, show_confidence=show_confidence)
|
|
137
|
+
return to_pil(annotated_bgr)
|
|
138
|
+
|
|
139
|
+
def annotate_and_save(
|
|
140
|
+
self,
|
|
141
|
+
image: ImageInput,
|
|
142
|
+
output_path: Union[str, Path],
|
|
143
|
+
conf_threshold: Optional[float] = None,
|
|
144
|
+
show_confidence: bool = True,
|
|
145
|
+
) -> Path:
|
|
146
|
+
"""Detect, annotate, and save result to *output_path*.
|
|
147
|
+
|
|
148
|
+
Returns the resolved output path.
|
|
149
|
+
"""
|
|
150
|
+
pil = self.annotate(image, conf_threshold=conf_threshold, show_confidence=show_confidence)
|
|
151
|
+
out = Path(output_path)
|
|
152
|
+
pil.save(out)
|
|
153
|
+
return out
|
|
154
|
+
|
|
155
|
+
# ------------------------------------------------------------------
|
|
156
|
+
# JSON export
|
|
157
|
+
# ------------------------------------------------------------------
|
|
158
|
+
|
|
159
|
+
def detect_to_json(
|
|
160
|
+
self,
|
|
161
|
+
image: ImageInput,
|
|
162
|
+
conf_threshold: Optional[float] = None,
|
|
163
|
+
indent: int = 2,
|
|
164
|
+
) -> str:
|
|
165
|
+
"""Return detection results as a JSON string."""
|
|
166
|
+
result = self.detect(image, conf_threshold=conf_threshold)
|
|
167
|
+
return json.dumps(result.to_dict(), indent=indent)
|
|
168
|
+
|
|
169
|
+
# ------------------------------------------------------------------
|
|
170
|
+
# Dunder
|
|
171
|
+
# ------------------------------------------------------------------
|
|
172
|
+
|
|
173
|
+
def __repr__(self) -> str:
|
|
174
|
+
return (
|
|
175
|
+
f"DentalDetector(model='{self.model_path.name}', "
|
|
176
|
+
f"conf_threshold={self.conf_threshold})"
|
|
177
|
+
)
|
|
@@ -0,0 +1,111 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from dataclasses import dataclass, field
|
|
4
|
+
from typing import List, Tuple
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
@dataclass
|
|
8
|
+
class Detection:
|
|
9
|
+
"""Single tooth detection result."""
|
|
10
|
+
|
|
11
|
+
label: str
|
|
12
|
+
class_id: int
|
|
13
|
+
confidence: float
|
|
14
|
+
bbox: Tuple[float, float, float, float] # (x1, y1, x2, y2) in pixels
|
|
15
|
+
|
|
16
|
+
@property
|
|
17
|
+
def center(self) -> Tuple[float, float]:
|
|
18
|
+
x1, y1, x2, y2 = self.bbox
|
|
19
|
+
return ((x1 + x2) / 2, (y1 + y2) / 2)
|
|
20
|
+
|
|
21
|
+
@property
|
|
22
|
+
def width(self) -> float:
|
|
23
|
+
return self.bbox[2] - self.bbox[0]
|
|
24
|
+
|
|
25
|
+
@property
|
|
26
|
+
def height(self) -> float:
|
|
27
|
+
return self.bbox[3] - self.bbox[1]
|
|
28
|
+
|
|
29
|
+
@property
|
|
30
|
+
def area(self) -> float:
|
|
31
|
+
return self.width * self.height
|
|
32
|
+
|
|
33
|
+
def to_dict(self) -> dict:
|
|
34
|
+
return {
|
|
35
|
+
"label": self.label,
|
|
36
|
+
"class_id": self.class_id,
|
|
37
|
+
"confidence": round(self.confidence, 4),
|
|
38
|
+
"bbox": {
|
|
39
|
+
"x1": round(self.bbox[0], 2),
|
|
40
|
+
"y1": round(self.bbox[1], 2),
|
|
41
|
+
"x2": round(self.bbox[2], 2),
|
|
42
|
+
"y2": round(self.bbox[3], 2),
|
|
43
|
+
},
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
@dataclass
|
|
48
|
+
class DetectionResult:
|
|
49
|
+
"""Full detection output for one image."""
|
|
50
|
+
|
|
51
|
+
detections: List[Detection] = field(default_factory=list)
|
|
52
|
+
image_width: int = 0
|
|
53
|
+
image_height: int = 0
|
|
54
|
+
|
|
55
|
+
@property
|
|
56
|
+
def labels(self) -> List[str]:
|
|
57
|
+
return [d.label for d in self.detections]
|
|
58
|
+
|
|
59
|
+
@property
|
|
60
|
+
def unique_labels(self) -> List[str]:
|
|
61
|
+
return sorted(set(self.labels))
|
|
62
|
+
|
|
63
|
+
def filter_by_label(self, label: str) -> "DetectionResult":
|
|
64
|
+
return DetectionResult(
|
|
65
|
+
detections=[d for d in self.detections if d.label == label],
|
|
66
|
+
image_width=self.image_width,
|
|
67
|
+
image_height=self.image_height,
|
|
68
|
+
)
|
|
69
|
+
|
|
70
|
+
def filter_by_confidence(self, threshold: float) -> "DetectionResult":
|
|
71
|
+
return DetectionResult(
|
|
72
|
+
detections=[d for d in self.detections if d.confidence >= threshold],
|
|
73
|
+
image_width=self.image_width,
|
|
74
|
+
image_height=self.image_height,
|
|
75
|
+
)
|
|
76
|
+
|
|
77
|
+
def to_dict(self) -> dict:
|
|
78
|
+
return {
|
|
79
|
+
"image_size": {"width": self.image_width, "height": self.image_height},
|
|
80
|
+
"count": len(self.detections),
|
|
81
|
+
"detections": [d.to_dict() for d in self.detections],
|
|
82
|
+
}
|
|
83
|
+
|
|
84
|
+
def __len__(self) -> int:
|
|
85
|
+
return len(self.detections)
|
|
86
|
+
|
|
87
|
+
def __iter__(self):
|
|
88
|
+
return iter(self.detections)
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
# Canonical tooth class names as trained
|
|
92
|
+
TOOTH_CLASSES = {
|
|
93
|
+
0: "1st Molar",
|
|
94
|
+
1: "1st Premolar",
|
|
95
|
+
2: "2nd Molar",
|
|
96
|
+
3: "2nd Premolar",
|
|
97
|
+
4: "Canine",
|
|
98
|
+
5: "Central Incisor",
|
|
99
|
+
6: "Lateral Incisor",
|
|
100
|
+
}
|
|
101
|
+
|
|
102
|
+
# FDI-style color map for visualization
|
|
103
|
+
LABEL_COLORS = {
|
|
104
|
+
"1st Molar": (255, 87, 51),
|
|
105
|
+
"2nd Molar": (255, 140, 0),
|
|
106
|
+
"1st Premolar": (50, 205, 50),
|
|
107
|
+
"2nd Premolar": (0, 180, 100),
|
|
108
|
+
"Canine": (30, 144, 255),
|
|
109
|
+
"Central Incisor": (148, 0, 211),
|
|
110
|
+
"Lateral Incisor": (255, 20, 147),
|
|
111
|
+
}
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import os
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
from typing import Union, Tuple
|
|
6
|
+
|
|
7
|
+
import cv2
|
|
8
|
+
import numpy as np
|
|
9
|
+
from PIL import Image
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
ImageInput = Union[str, Path, np.ndarray, Image.Image]
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def load_image(source: ImageInput) -> np.ndarray:
|
|
16
|
+
"""Load any supported image type into a BGR numpy array."""
|
|
17
|
+
if isinstance(source, (str, Path)):
|
|
18
|
+
path = str(source)
|
|
19
|
+
if not os.path.exists(path):
|
|
20
|
+
raise FileNotFoundError(f"Image not found: {path}")
|
|
21
|
+
img = cv2.imread(path)
|
|
22
|
+
if img is None:
|
|
23
|
+
raise ValueError(f"Could not decode image: {path}")
|
|
24
|
+
return img
|
|
25
|
+
if isinstance(source, Image.Image):
|
|
26
|
+
return cv2.cvtColor(np.array(source.convert("RGB")), cv2.COLOR_RGB2BGR)
|
|
27
|
+
if isinstance(source, np.ndarray):
|
|
28
|
+
if source.ndim == 2:
|
|
29
|
+
return cv2.cvtColor(source, cv2.COLOR_GRAY2BGR)
|
|
30
|
+
if source.shape[2] == 4:
|
|
31
|
+
return cv2.cvtColor(source, cv2.COLOR_RGBA2BGR)
|
|
32
|
+
if source.shape[2] == 3:
|
|
33
|
+
# Assume RGB from PIL/matplotlib; convert to BGR for OpenCV
|
|
34
|
+
return source[:, :, ::-1].copy()
|
|
35
|
+
raise ValueError(f"Unsupported array shape: {source.shape}")
|
|
36
|
+
raise TypeError(f"Unsupported image type: {type(source)}")
|
|
37
|
+
|
|
38
|
+
|
|
39
|
+
def to_pil(image: np.ndarray) -> Image.Image:
|
|
40
|
+
"""Convert a BGR numpy array to a PIL RGB Image."""
|
|
41
|
+
return Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def letterbox(
|
|
45
|
+
image: np.ndarray,
|
|
46
|
+
target: Tuple[int, int] = (640, 640),
|
|
47
|
+
pad_value: int = 114,
|
|
48
|
+
) -> Tuple[np.ndarray, float, Tuple[int, int]]:
|
|
49
|
+
"""Resize with aspect ratio preserved and pad to `target` (W, H).
|
|
50
|
+
|
|
51
|
+
Returns:
|
|
52
|
+
padded image, scale factor, (left_pad, top_pad)
|
|
53
|
+
"""
|
|
54
|
+
h, w = image.shape[:2]
|
|
55
|
+
tw, th = target
|
|
56
|
+
scale = min(tw / w, th / h)
|
|
57
|
+
nw, nh = int(w * scale), int(h * scale)
|
|
58
|
+
resized = cv2.resize(image, (nw, nh), interpolation=cv2.INTER_LINEAR)
|
|
59
|
+
canvas = np.full((th, tw, 3), pad_value, dtype=np.uint8)
|
|
60
|
+
left, top = (tw - nw) // 2, (th - nh) // 2
|
|
61
|
+
canvas[top : top + nh, left : left + nw] = resized
|
|
62
|
+
return canvas, scale, (left, top)
|
|
@@ -0,0 +1,61 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from typing import Optional
|
|
4
|
+
|
|
5
|
+
import cv2
|
|
6
|
+
import numpy as np
|
|
7
|
+
|
|
8
|
+
from .models import Detection, DetectionResult, LABEL_COLORS
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
_DEFAULT_COLOR = (0, 255, 0)
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def annotate(
|
|
15
|
+
image: np.ndarray,
|
|
16
|
+
result: DetectionResult,
|
|
17
|
+
show_confidence: bool = True,
|
|
18
|
+
line_thickness: int = 2,
|
|
19
|
+
font_scale: float = 0.55,
|
|
20
|
+
) -> np.ndarray:
|
|
21
|
+
"""Draw bounding boxes and labels on a BGR image copy."""
|
|
22
|
+
out = image.copy()
|
|
23
|
+
for det in result.detections:
|
|
24
|
+
color = LABEL_COLORS.get(det.label, _DEFAULT_COLOR)
|
|
25
|
+
# BGR order for OpenCV
|
|
26
|
+
bgr = (color[2], color[1], color[0])
|
|
27
|
+
x1, y1, x2, y2 = (int(v) for v in det.bbox)
|
|
28
|
+
cv2.rectangle(out, (x1, y1), (x2, y2), bgr, line_thickness)
|
|
29
|
+
label_text = f"{det.label} {det.confidence:.2f}" if show_confidence else det.label
|
|
30
|
+
_draw_label(out, label_text, x1, y1, bgr, font_scale, line_thickness)
|
|
31
|
+
return out
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _draw_label(
|
|
35
|
+
image: np.ndarray,
|
|
36
|
+
text: str,
|
|
37
|
+
x: int,
|
|
38
|
+
y: int,
|
|
39
|
+
color,
|
|
40
|
+
font_scale: float,
|
|
41
|
+
thickness: int,
|
|
42
|
+
) -> None:
|
|
43
|
+
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
44
|
+
(tw, th), baseline = cv2.getTextSize(text, font, font_scale, thickness)
|
|
45
|
+
# Background rectangle above the bounding box
|
|
46
|
+
bg_y1 = max(0, y - th - baseline - 4)
|
|
47
|
+
bg_y2 = max(th + baseline + 4, y)
|
|
48
|
+
cv2.rectangle(image, (x, bg_y1), (x + tw + 4, bg_y2), color, -1)
|
|
49
|
+
text_color = _contrast_color(color)
|
|
50
|
+
cv2.putText(
|
|
51
|
+
image, text,
|
|
52
|
+
(x + 2, bg_y2 - baseline - 2),
|
|
53
|
+
font, font_scale, text_color, thickness, cv2.LINE_AA,
|
|
54
|
+
)
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def _contrast_color(bgr) -> tuple:
|
|
58
|
+
"""Return black or white depending on perceived luminance."""
|
|
59
|
+
b, g, r = bgr
|
|
60
|
+
luminance = 0.299 * r + 0.587 * g + 0.114 * b
|
|
61
|
+
return (0, 0, 0) if luminance > 128 else (255, 255, 255)
|
|
@@ -0,0 +1,77 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "dental-detector"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Tooth detection in maxilla and mandible dental images using YOLOv8"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = { text = "MIT" }
|
|
11
|
+
requires-python = ">=3.8"
|
|
12
|
+
authors = [
|
|
13
|
+
{ name = "Sandeep Vissa", email = "sandeepvissa2002@gmail.com" },
|
|
14
|
+
]
|
|
15
|
+
keywords = [
|
|
16
|
+
"dental",
|
|
17
|
+
"teeth",
|
|
18
|
+
"tooth-detection",
|
|
19
|
+
"yolo",
|
|
20
|
+
"object-detection",
|
|
21
|
+
"maxilla",
|
|
22
|
+
"mandible",
|
|
23
|
+
"dental-imaging",
|
|
24
|
+
"computer-vision",
|
|
25
|
+
]
|
|
26
|
+
classifiers = [
|
|
27
|
+
"Development Status :: 3 - Alpha",
|
|
28
|
+
"Intended Audience :: Developers",
|
|
29
|
+
"Intended Audience :: Science/Research",
|
|
30
|
+
"Intended Audience :: Healthcare Industry",
|
|
31
|
+
"License :: OSI Approved :: MIT License",
|
|
32
|
+
"Operating System :: OS Independent",
|
|
33
|
+
"Programming Language :: Python :: 3",
|
|
34
|
+
"Programming Language :: Python :: 3.8",
|
|
35
|
+
"Programming Language :: Python :: 3.9",
|
|
36
|
+
"Programming Language :: Python :: 3.10",
|
|
37
|
+
"Programming Language :: Python :: 3.11",
|
|
38
|
+
"Programming Language :: Python :: 3.12",
|
|
39
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
40
|
+
"Topic :: Scientific/Engineering :: Image Recognition",
|
|
41
|
+
"Topic :: Scientific/Engineering :: Medical Science Apps.",
|
|
42
|
+
]
|
|
43
|
+
dependencies = [
|
|
44
|
+
"ultralytics>=8.0.0",
|
|
45
|
+
"opencv-python>=4.5.0",
|
|
46
|
+
"Pillow>=9.0.0",
|
|
47
|
+
"numpy>=1.21.0",
|
|
48
|
+
]
|
|
49
|
+
|
|
50
|
+
[project.optional-dependencies]
|
|
51
|
+
dev = [
|
|
52
|
+
"pytest>=7.0",
|
|
53
|
+
"pytest-cov>=4.0",
|
|
54
|
+
"ruff>=0.1.0",
|
|
55
|
+
]
|
|
56
|
+
|
|
57
|
+
[project.urls]
|
|
58
|
+
Homepage = "https://github.com/sandeepvissa/dental-detector"
|
|
59
|
+
Repository = "https://github.com/sandeepvissa/dental-detector"
|
|
60
|
+
"Bug Tracker" = "https://github.com/sandeepvissa/dental-detector/issues"
|
|
61
|
+
|
|
62
|
+
[project.scripts]
|
|
63
|
+
dental-detector = "dental_detector.cli:main"
|
|
64
|
+
|
|
65
|
+
[tool.hatch.build.targets.wheel]
|
|
66
|
+
packages = ["dental_detector"]
|
|
67
|
+
|
|
68
|
+
[tool.ruff]
|
|
69
|
+
line-length = 100
|
|
70
|
+
target-version = "py38"
|
|
71
|
+
|
|
72
|
+
[tool.ruff.lint]
|
|
73
|
+
select = ["E", "F", "W", "I"]
|
|
74
|
+
|
|
75
|
+
[tool.pytest.ini_options]
|
|
76
|
+
testpaths = ["tests"]
|
|
77
|
+
addopts = "-v --tb=short"
|
|
File without changes
|
|
@@ -0,0 +1,344 @@
|
|
|
1
|
+
"""Tests for dental_detector package."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
from unittest.mock import MagicMock, patch
|
|
8
|
+
|
|
9
|
+
import numpy as np
|
|
10
|
+
import pytest
|
|
11
|
+
from PIL import Image
|
|
12
|
+
|
|
13
|
+
from dental_detector import DentalDetector, Detection, DetectionResult, TOOTH_CLASSES
|
|
14
|
+
from dental_detector.utils import load_image, letterbox, to_pil
|
|
15
|
+
from dental_detector.visualization import annotate
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
# ---------------------------------------------------------------------------
|
|
19
|
+
# Fixtures
|
|
20
|
+
# ---------------------------------------------------------------------------
|
|
21
|
+
|
|
22
|
+
@pytest.fixture
|
|
23
|
+
def dummy_image_rgb() -> np.ndarray:
|
|
24
|
+
"""100x80 white RGB image."""
|
|
25
|
+
return np.full((80, 100, 3), 255, dtype=np.uint8)
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
@pytest.fixture
|
|
29
|
+
def dummy_pil_image() -> Image.Image:
|
|
30
|
+
return Image.fromarray(np.zeros((64, 64, 3), dtype=np.uint8))
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@pytest.fixture
|
|
34
|
+
def sample_detection() -> Detection:
|
|
35
|
+
return Detection(
|
|
36
|
+
label="Central Incisor",
|
|
37
|
+
class_id=5,
|
|
38
|
+
confidence=0.92,
|
|
39
|
+
bbox=(10.0, 20.0, 80.0, 70.0),
|
|
40
|
+
)
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
@pytest.fixture
|
|
44
|
+
def sample_result(sample_detection) -> DetectionResult:
|
|
45
|
+
return DetectionResult(
|
|
46
|
+
detections=[sample_detection],
|
|
47
|
+
image_width=100,
|
|
48
|
+
image_height=80,
|
|
49
|
+
)
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
# ---------------------------------------------------------------------------
|
|
53
|
+
# Detection dataclass
|
|
54
|
+
# ---------------------------------------------------------------------------
|
|
55
|
+
|
|
56
|
+
class TestDetection:
|
|
57
|
+
def test_center(self, sample_detection):
|
|
58
|
+
cx, cy = sample_detection.center
|
|
59
|
+
assert cx == pytest.approx(45.0)
|
|
60
|
+
assert cy == pytest.approx(45.0)
|
|
61
|
+
|
|
62
|
+
def test_width_height(self, sample_detection):
|
|
63
|
+
assert sample_detection.width == pytest.approx(70.0)
|
|
64
|
+
assert sample_detection.height == pytest.approx(50.0)
|
|
65
|
+
|
|
66
|
+
def test_area(self, sample_detection):
|
|
67
|
+
assert sample_detection.area == pytest.approx(3500.0)
|
|
68
|
+
|
|
69
|
+
def test_to_dict_keys(self, sample_detection):
|
|
70
|
+
d = sample_detection.to_dict()
|
|
71
|
+
assert set(d.keys()) == {"label", "class_id", "confidence", "bbox"}
|
|
72
|
+
assert d["label"] == "Central Incisor"
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
# ---------------------------------------------------------------------------
|
|
76
|
+
# DetectionResult
|
|
77
|
+
# ---------------------------------------------------------------------------
|
|
78
|
+
|
|
79
|
+
class TestDetectionResult:
|
|
80
|
+
def test_len(self, sample_result):
|
|
81
|
+
assert len(sample_result) == 1
|
|
82
|
+
|
|
83
|
+
def test_labels(self, sample_result):
|
|
84
|
+
assert sample_result.labels == ["Central Incisor"]
|
|
85
|
+
|
|
86
|
+
def test_unique_labels(self):
|
|
87
|
+
r = DetectionResult(
|
|
88
|
+
detections=[
|
|
89
|
+
Detection("Canine", 4, 0.8, (0, 0, 10, 10)),
|
|
90
|
+
Detection("Canine", 4, 0.7, (10, 10, 20, 20)),
|
|
91
|
+
Detection("1st Molar", 0, 0.9, (20, 20, 30, 30)),
|
|
92
|
+
],
|
|
93
|
+
)
|
|
94
|
+
assert r.unique_labels == ["1st Molar", "Canine"]
|
|
95
|
+
|
|
96
|
+
def test_filter_by_label(self, sample_result):
|
|
97
|
+
filtered = sample_result.filter_by_label("Central Incisor")
|
|
98
|
+
assert len(filtered) == 1
|
|
99
|
+
filtered_empty = sample_result.filter_by_label("Canine")
|
|
100
|
+
assert len(filtered_empty) == 0
|
|
101
|
+
|
|
102
|
+
def test_filter_by_confidence(self):
|
|
103
|
+
r = DetectionResult(
|
|
104
|
+
detections=[
|
|
105
|
+
Detection("Canine", 4, 0.9, (0, 0, 10, 10)),
|
|
106
|
+
Detection("1st Molar", 0, 0.4, (20, 20, 30, 30)),
|
|
107
|
+
],
|
|
108
|
+
)
|
|
109
|
+
high = r.filter_by_confidence(0.6)
|
|
110
|
+
assert len(high) == 1
|
|
111
|
+
assert high.detections[0].label == "Canine"
|
|
112
|
+
|
|
113
|
+
def test_to_dict_structure(self, sample_result):
|
|
114
|
+
d = sample_result.to_dict()
|
|
115
|
+
assert "image_size" in d
|
|
116
|
+
assert "count" in d
|
|
117
|
+
assert "detections" in d
|
|
118
|
+
assert d["count"] == 1
|
|
119
|
+
|
|
120
|
+
def test_iteration(self, sample_result):
|
|
121
|
+
items = list(sample_result)
|
|
122
|
+
assert len(items) == 1
|
|
123
|
+
assert items[0].label == "Central Incisor"
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
# ---------------------------------------------------------------------------
|
|
127
|
+
# TOOTH_CLASSES
|
|
128
|
+
# ---------------------------------------------------------------------------
|
|
129
|
+
|
|
130
|
+
def test_tooth_classes_completeness():
|
|
131
|
+
assert len(TOOTH_CLASSES) == 7
|
|
132
|
+
assert 0 in TOOTH_CLASSES and TOOTH_CLASSES[0] == "1st Molar"
|
|
133
|
+
assert 5 in TOOTH_CLASSES and TOOTH_CLASSES[5] == "Central Incisor"
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
# ---------------------------------------------------------------------------
|
|
137
|
+
# Utils
|
|
138
|
+
# ---------------------------------------------------------------------------
|
|
139
|
+
|
|
140
|
+
class TestLoadImage:
|
|
141
|
+
def test_load_from_numpy_rgb(self, dummy_image_rgb):
|
|
142
|
+
bgr = load_image(dummy_image_rgb)
|
|
143
|
+
assert bgr.shape == (80, 100, 3)
|
|
144
|
+
|
|
145
|
+
def test_load_from_pil(self, dummy_pil_image):
|
|
146
|
+
bgr = load_image(dummy_pil_image)
|
|
147
|
+
assert bgr.ndim == 3
|
|
148
|
+
assert bgr.shape[2] == 3
|
|
149
|
+
|
|
150
|
+
def test_load_from_path(self, tmp_path, dummy_pil_image):
|
|
151
|
+
p = tmp_path / "test.jpg"
|
|
152
|
+
dummy_pil_image.save(p)
|
|
153
|
+
bgr = load_image(p)
|
|
154
|
+
assert bgr.ndim == 3
|
|
155
|
+
|
|
156
|
+
def test_missing_file_raises(self):
|
|
157
|
+
with pytest.raises(FileNotFoundError):
|
|
158
|
+
load_image("/nonexistent/image.jpg")
|
|
159
|
+
|
|
160
|
+
def test_wrong_type_raises(self):
|
|
161
|
+
with pytest.raises(TypeError):
|
|
162
|
+
load_image(12345)
|
|
163
|
+
|
|
164
|
+
|
|
165
|
+
class TestLetterbox:
|
|
166
|
+
def test_output_shape(self, dummy_image_rgb):
|
|
167
|
+
import cv2
|
|
168
|
+
bgr = cv2.cvtColor(dummy_image_rgb, cv2.COLOR_RGB2BGR)
|
|
169
|
+
out, scale, (left, top) = letterbox(bgr, (640, 640))
|
|
170
|
+
assert out.shape == (640, 640, 3)
|
|
171
|
+
|
|
172
|
+
def test_aspect_ratio_preserved(self):
|
|
173
|
+
import cv2
|
|
174
|
+
img = np.zeros((200, 100, 3), dtype=np.uint8)
|
|
175
|
+
out, scale, (left, top) = letterbox(img, (640, 640))
|
|
176
|
+
assert out.shape == (640, 640, 3)
|
|
177
|
+
assert scale == pytest.approx(3.2)
|
|
178
|
+
|
|
179
|
+
|
|
180
|
+
class TestToPil:
|
|
181
|
+
def test_returns_pil_image(self, dummy_image_rgb):
|
|
182
|
+
import cv2
|
|
183
|
+
bgr = cv2.cvtColor(dummy_image_rgb, cv2.COLOR_RGB2BGR)
|
|
184
|
+
pil = to_pil(bgr)
|
|
185
|
+
assert isinstance(pil, Image.Image)
|
|
186
|
+
assert pil.mode == "RGB"
|
|
187
|
+
|
|
188
|
+
|
|
189
|
+
# ---------------------------------------------------------------------------
|
|
190
|
+
# Visualization
|
|
191
|
+
# ---------------------------------------------------------------------------
|
|
192
|
+
|
|
193
|
+
class TestAnnotate:
|
|
194
|
+
def test_returns_same_shape(self, dummy_image_rgb, sample_result):
|
|
195
|
+
import cv2
|
|
196
|
+
bgr = cv2.cvtColor(dummy_image_rgb, cv2.COLOR_RGB2BGR)
|
|
197
|
+
out = annotate(bgr, sample_result)
|
|
198
|
+
assert out.shape == bgr.shape
|
|
199
|
+
|
|
200
|
+
def test_does_not_mutate_input(self, dummy_image_rgb, sample_result):
|
|
201
|
+
import cv2
|
|
202
|
+
bgr = cv2.cvtColor(dummy_image_rgb, cv2.COLOR_RGB2BGR)
|
|
203
|
+
original = bgr.copy()
|
|
204
|
+
annotate(bgr, sample_result)
|
|
205
|
+
np.testing.assert_array_equal(bgr, original)
|
|
206
|
+
|
|
207
|
+
def test_empty_result_unchanged(self, dummy_image_rgb):
|
|
208
|
+
import cv2
|
|
209
|
+
bgr = cv2.cvtColor(dummy_image_rgb, cv2.COLOR_RGB2BGR)
|
|
210
|
+
empty = DetectionResult(detections=[], image_width=100, image_height=80)
|
|
211
|
+
out = annotate(bgr, empty)
|
|
212
|
+
np.testing.assert_array_equal(out, bgr)
|
|
213
|
+
|
|
214
|
+
|
|
215
|
+
# ---------------------------------------------------------------------------
|
|
216
|
+
# DentalDetector (unit — YOLO mocked)
|
|
217
|
+
# ---------------------------------------------------------------------------
|
|
218
|
+
|
|
219
|
+
def _make_mock_model(detections):
|
|
220
|
+
"""Build a mock ultralytics YOLO result."""
|
|
221
|
+
import torch
|
|
222
|
+
|
|
223
|
+
boxes_mock = MagicMock()
|
|
224
|
+
boxes_mock.xyxy.cpu().numpy.return_value = np.array(
|
|
225
|
+
[[d["bbox"]] for d in detections], dtype=np.float32
|
|
226
|
+
).reshape(-1, 4) if detections else np.empty((0, 4), dtype=np.float32)
|
|
227
|
+
boxes_mock.conf.cpu().numpy.return_value = np.array(
|
|
228
|
+
[d["conf"] for d in detections], dtype=np.float32
|
|
229
|
+
)
|
|
230
|
+
boxes_mock.cls.cpu().numpy.return_value = np.array(
|
|
231
|
+
[d["cls"] for d in detections], dtype=np.float32
|
|
232
|
+
)
|
|
233
|
+
|
|
234
|
+
result_mock = MagicMock()
|
|
235
|
+
result_mock.boxes = boxes_mock
|
|
236
|
+
result_mock.names = {int(k): v for k, v in TOOTH_CLASSES.items()}
|
|
237
|
+
|
|
238
|
+
model_mock = MagicMock()
|
|
239
|
+
model_mock.return_value = [result_mock]
|
|
240
|
+
return model_mock
|
|
241
|
+
|
|
242
|
+
|
|
243
|
+
@pytest.fixture
|
|
244
|
+
def mock_detector(tmp_path):
|
|
245
|
+
"""DentalDetector with a dummy weights file and mocked YOLO."""
|
|
246
|
+
weights = tmp_path / "dummy.pt"
|
|
247
|
+
weights.touch()
|
|
248
|
+
|
|
249
|
+
with patch("dental_detector.detector.YOLO") as MockYOLO:
|
|
250
|
+
mock_model = _make_mock_model([
|
|
251
|
+
{"bbox": (5, 10, 50, 60), "conf": 0.88, "cls": 5},
|
|
252
|
+
{"bbox": (60, 10, 95, 60), "conf": 0.55, "cls": 4},
|
|
253
|
+
])
|
|
254
|
+
MockYOLO.return_value = mock_model
|
|
255
|
+
detector = DentalDetector(str(weights), conf_threshold=0.6)
|
|
256
|
+
detector._model = mock_model
|
|
257
|
+
yield detector
|
|
258
|
+
|
|
259
|
+
|
|
260
|
+
class TestDentalDetector:
|
|
261
|
+
def test_detect_returns_result(self, mock_detector, dummy_image_rgb):
|
|
262
|
+
result = mock_detector.detect(dummy_image_rgb)
|
|
263
|
+
assert isinstance(result, DetectionResult)
|
|
264
|
+
|
|
265
|
+
def test_confidence_filtering(self, mock_detector, dummy_image_rgb):
|
|
266
|
+
result = mock_detector.detect(dummy_image_rgb, conf_threshold=0.6)
|
|
267
|
+
for det in result:
|
|
268
|
+
assert det.confidence >= 0.6
|
|
269
|
+
|
|
270
|
+
def test_detect_to_json(self, mock_detector, dummy_image_rgb):
|
|
271
|
+
js = mock_detector.detect_to_json(dummy_image_rgb)
|
|
272
|
+
data = json.loads(js)
|
|
273
|
+
assert "detections" in data
|
|
274
|
+
assert "count" in data
|
|
275
|
+
|
|
276
|
+
def test_annotate_returns_pil(self, mock_detector, dummy_image_rgb):
|
|
277
|
+
pil = mock_detector.annotate(dummy_image_rgb)
|
|
278
|
+
assert isinstance(pil, Image.Image)
|
|
279
|
+
|
|
280
|
+
def test_annotate_and_save(self, mock_detector, dummy_image_rgb, tmp_path):
|
|
281
|
+
out = tmp_path / "out.jpg"
|
|
282
|
+
returned = mock_detector.annotate_and_save(dummy_image_rgb, out)
|
|
283
|
+
assert returned == out
|
|
284
|
+
assert out.exists()
|
|
285
|
+
|
|
286
|
+
def test_model_not_found_raises(self):
|
|
287
|
+
with pytest.raises(FileNotFoundError):
|
|
288
|
+
DentalDetector("/nonexistent/model.pt")
|
|
289
|
+
|
|
290
|
+
def test_repr(self, mock_detector):
|
|
291
|
+
r = repr(mock_detector)
|
|
292
|
+
assert "DentalDetector" in r
|
|
293
|
+
assert "conf_threshold" in r
|
|
294
|
+
|
|
295
|
+
|
|
296
|
+
# ---------------------------------------------------------------------------
|
|
297
|
+
# CLI
|
|
298
|
+
# ---------------------------------------------------------------------------
|
|
299
|
+
|
|
300
|
+
class TestCLI:
|
|
301
|
+
def test_detect_prints_output(self, mock_detector, dummy_pil_image, tmp_path, capsys):
|
|
302
|
+
img_path = tmp_path / "test.jpg"
|
|
303
|
+
dummy_pil_image.save(img_path)
|
|
304
|
+
weights = tmp_path / "dummy.pt"
|
|
305
|
+
weights.touch()
|
|
306
|
+
|
|
307
|
+
from dental_detector.cli import main
|
|
308
|
+
|
|
309
|
+
with patch("dental_detector.cli.DentalDetector") as MockDet:
|
|
310
|
+
inst = MagicMock()
|
|
311
|
+
inst.detect.return_value = DetectionResult(
|
|
312
|
+
detections=[Detection("1st Molar", 0, 0.9, (0, 0, 10, 10))],
|
|
313
|
+
image_width=64,
|
|
314
|
+
image_height=64,
|
|
315
|
+
)
|
|
316
|
+
MockDet.return_value = inst
|
|
317
|
+
rc = main([str(img_path), "--model", str(weights)])
|
|
318
|
+
|
|
319
|
+
assert rc == 0
|
|
320
|
+
out = capsys.readouterr().out
|
|
321
|
+
assert "1st Molar" in out
|
|
322
|
+
|
|
323
|
+
def test_json_flag(self, dummy_pil_image, tmp_path, capsys):
|
|
324
|
+
img_path = tmp_path / "test.jpg"
|
|
325
|
+
dummy_pil_image.save(img_path)
|
|
326
|
+
weights = tmp_path / "dummy.pt"
|
|
327
|
+
weights.touch()
|
|
328
|
+
|
|
329
|
+
from dental_detector.cli import main
|
|
330
|
+
|
|
331
|
+
with patch("dental_detector.cli.DentalDetector") as MockDet:
|
|
332
|
+
inst = MagicMock()
|
|
333
|
+
inst.detect.return_value = DetectionResult(
|
|
334
|
+
detections=[Detection("Canine", 4, 0.75, (5, 5, 20, 20))],
|
|
335
|
+
image_width=64,
|
|
336
|
+
image_height=64,
|
|
337
|
+
)
|
|
338
|
+
MockDet.return_value = inst
|
|
339
|
+
rc = main([str(img_path), "--model", str(weights), "--json"])
|
|
340
|
+
|
|
341
|
+
assert rc == 0
|
|
342
|
+
out = capsys.readouterr().out
|
|
343
|
+
data = json.loads(out)
|
|
344
|
+
assert "detections" in data
|