VietCat commited on
Commit
6998ed6
·
1 Parent(s): ef364be
Files changed (1) hide show
  1. model.py +15 -2
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
- self.box_color = config['inference']['box_color']
38
- self.text_color = config['inference']['text_color']
 
 
 
 
 
 
 
 
 
 
 
 
 
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