Update model_1.py

model_1.py

@@ -42,6 +42,12 @@ if 'StatefulPartitionedCall' in outname:
 else:
     boxes_idx, classes_idx, scores_idx = 0, 1, 2
 
+# Define the color mapping
+color_mapping = {
+    'Misalignment class': (0, 0, 255),  # Red for misalignment
+    'Empty class': (0, 255, 0)          # Green for empty
+}
+
 def perform_detection(image):
     imH, imW, _ = image.shape
     image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
@@ -66,8 +72,10 @@ def perform_detection(image):
             ymax = int(min(imH, (boxes[i][2] * imH)))
             xmax = int(min(imW, (boxes[i][3] * imW)))
 
-            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), (10, 255, 0), 2)
             object_name = labels[int(classes[i])]
+            color = color_mapping.get(object_name, (10, 255, 0))  # Default color if not specified
+
+            cv2.rectangle(image, (xmin, ymin), (xmax, ymax), color, 2)
             label = '%s: %d%%' % (object_name, int(scores[i] * 100))
             labelSize, baseLine = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
             label_ymin = max(ymin, labelSize[1] + 10)