BenjiELCA commited on
Commit
7440d31
1 Parent(s): 54c08aa

no subprocess problem

Browse files
Files changed (1) hide show
  1. modules/eval.py +5 -0
modules/eval.py CHANGED
@@ -101,6 +101,11 @@ def object_prediction(model, image, score_threshold=0.5, iou_threshold=0.5):
101
  scores = scores[selected_boxes]
102
  labels = labels[selected_boxes]
103
 
 
 
 
 
 
104
  prediction = {
105
  'boxes': boxes,
106
  'scores': scores,
 
101
  scores = scores[selected_boxes]
102
  labels = labels[selected_boxes]
103
 
104
+ #modify the label of the sub-process to task
105
+ for i in range(len(labels)):
106
+ if labels[i] == list(object_dict.values()).index('subProcess'):
107
+ labels[i] = list(object_dict.values()).index('task')
108
+
109
  prediction = {
110
  'boxes': boxes,
111
  'scores': scores,