dtyago commited on
Commit
e3f36c1
1 Parent(s): bf66ad3

colors not supported so refactored

Browse files
Files changed (1) hide show
  1. app.py +6 -8
app.py CHANGED
@@ -35,9 +35,9 @@ def predict_combined_models(image_np, model1, model2):
35
  results2 = model2(image_np)
36
  result2 = {key:value.numpy() for key,value in results2.items()}
37
 
38
- # Define colors for bounding boxes for each model
39
- color_final = (0, 255, 0) # Green for model 1
40
- color_initial = (255, 0, 0) # Red for model 2
41
 
42
  # Visualization for model 1
43
  image_np_with_detections = image_np.copy()
@@ -46,26 +46,24 @@ def predict_combined_models(image_np, model1, model2):
46
  result1['detection_boxes'][0],
47
  (result1['detection_classes'][0]).astype(int),
48
  result1['detection_scores'][0],
49
- category_index,
50
  use_normalized_coordinates=True,
51
  max_boxes_to_draw=200,
52
  min_score_thresh=.60,
53
  agnostic_mode=False,
54
- colors=color_final,
55
  line_thickness=2)
56
 
57
- # Visualization for model 2 (can adjust styles to differentiate)
58
  viz_utils.visualize_boxes_and_labels_on_image_array(
59
  image_np_with_detections[0],
60
  result2['detection_boxes'][0],
61
  (result2['detection_classes'][0]).astype(int),
62
  result2['detection_scores'][0],
63
- category_index,
64
  use_normalized_coordinates=True,
65
  max_boxes_to_draw=200,
66
  min_score_thresh=.60,
67
  agnostic_mode=False,
68
- colors=color_initial,
69
  line_thickness=2)
70
 
71
  # Combine and return final image
 
35
  results2 = model2(image_np)
36
  result2 = {key:value.numpy() for key,value in results2.items()}
37
 
38
+ # Modify category index for each model
39
+ category_index_model1 = {k: {**v, 'name': v['name'] + ' - Model 1'} for k, v in category_index.items()}
40
+ category_index_model2 = {k: {**v, 'name': v['name'] + ' - Model 2'} for k, v in category_index.items()}
41
 
42
  # Visualization for model 1
43
  image_np_with_detections = image_np.copy()
 
46
  result1['detection_boxes'][0],
47
  (result1['detection_classes'][0]).astype(int),
48
  result1['detection_scores'][0],
49
+ category_index_model1,
50
  use_normalized_coordinates=True,
51
  max_boxes_to_draw=200,
52
  min_score_thresh=.60,
53
  agnostic_mode=False,
 
54
  line_thickness=2)
55
 
56
+ # Visualization for model 2
57
  viz_utils.visualize_boxes_and_labels_on_image_array(
58
  image_np_with_detections[0],
59
  result2['detection_boxes'][0],
60
  (result2['detection_classes'][0]).astype(int),
61
  result2['detection_scores'][0],
62
+ category_index_model2,
63
  use_normalized_coordinates=True,
64
  max_boxes_to_draw=200,
65
  min_score_thresh=.60,
66
  agnostic_mode=False,
 
67
  line_thickness=2)
68
 
69
  # Combine and return final image