sshi commited on
Commit
f6b2407
1 Parent(s): 4a0733d

App bug fix.

Browse files
Files changed (1) hide show
  1. app.py +4 -7
app.py CHANGED
@@ -82,11 +82,9 @@ model_yolos.to(device)
82
  model_yolos.eval()
83
 
84
  # colors for visualization
85
- COLORS = [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],
86
- [0.756, 0.794, 0.100], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933],
87
- [0.184, 0.494, 0.741], [0.494, 0.674, 0.556], [0.494, 0.301, 0.933],
88
- [0.000, 0.325, 0.850], [0.745, 0.301, 0.188]]
89
-
90
  # for output bounding box post-processing
91
  def box_cxcywh_to_xyxy(x):
92
  x_c, y_c, w, h = x.unbind(1)
@@ -103,11 +101,10 @@ def rescale_bboxes(out_bbox, size):
103
  def plot_results(pil_img, prob, boxes):
104
 
105
  draw = ImageDraw.Draw(pil_img)
106
- colors = COLORS * 100
107
 
108
  for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
109
  cl = p.argmax()
110
- c = tuple(colors[cl])
111
 
112
  draw.rectangle([xmin, ymin, xmax - xmin, ymax - ymin], outline=c, width=2)
113
  draw.text(
 
82
  model_yolos.eval()
83
 
84
  # colors for visualization
85
+ colors = ['Pink', 'Crimson', 'Yellow', 'Indigo', 'BlueViolet',
86
+ 'Blue', 'GhostWhite', 'LightSteelBlue', 'Brown', 'SkyBlue',
87
+ 'Tomato']
 
 
88
  # for output bounding box post-processing
89
  def box_cxcywh_to_xyxy(x):
90
  x_c, y_c, w, h = x.unbind(1)
 
101
  def plot_results(pil_img, prob, boxes):
102
 
103
  draw = ImageDraw.Draw(pil_img)
 
104
 
105
  for p, (xmin, ymin, xmax, ymax) in zip(prob, boxes.tolist()):
106
  cl = p.argmax()
107
+ c = colors[cl]
108
 
109
  draw.rectangle([xmin, ymin, xmax - xmin, ymax - ymin], outline=c, width=2)
110
  draw.text(