Spaces:
Sleeping
Sleeping
fix bug
Browse files
model.py
CHANGED
|
@@ -34,8 +34,21 @@ class TrafficSignDetector:
|
|
| 34 |
torch.load = original_torch_load
|
| 35 |
|
| 36 |
self.conf_threshold = config['model']['confidence_threshold']
|
| 37 |
-
|
| 38 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 39 |
self.thickness = config['inference']['thickness']
|
| 40 |
self.classes = config['classes']
|
| 41 |
|
|
|
|
| 34 |
torch.load = original_torch_load
|
| 35 |
|
| 36 |
self.conf_threshold = config['model']['confidence_threshold']
|
| 37 |
+
|
| 38 |
+
# Convert color strings to tuples if needed
|
| 39 |
+
box_color = config['inference']['box_color']
|
| 40 |
+
if isinstance(box_color, str):
|
| 41 |
+
# Convert string "(128, 0, 128)" to tuple (128, 0, 128)
|
| 42 |
+
self.box_color = tuple(map(int, box_color.strip('()').split(',')))
|
| 43 |
+
else:
|
| 44 |
+
self.box_color = box_color
|
| 45 |
+
|
| 46 |
+
text_color = config['inference']['text_color']
|
| 47 |
+
if isinstance(text_color, str):
|
| 48 |
+
self.text_color = tuple(map(int, text_color.strip('()').split(',')))
|
| 49 |
+
else:
|
| 50 |
+
self.text_color = text_color
|
| 51 |
+
|
| 52 |
self.thickness = config['inference']['thickness']
|
| 53 |
self.classes = config['classes']
|
| 54 |
|