Update demo.py

demo.py

@@ -53,7 +53,11 @@ def fire(frame):
         # image = cv2.resize(image, (1024, 640))
         for obj_name in object_name:
             # print(plate_number)
-            image = plot_one_box(box[obj_name], image, label=object_name[obj_name], score=pro[obj_name], color=[0, 255, 0], line_thickness=1)
+            if object_name[obj_name]:
+                if object_name[obj_name] == "fire":
+                    image = plot_one_box(box[obj_name], image, label=object_name[obj_name], score=pro[obj_name], color=[0, 255, 0], line_thickness=1)
+                else:
+                    image = plot_one_box(box[obj_name], image, label=object_name[obj_name], score=pro[obj_name], color=[0, 0, 255], line_thickness=1)
         image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
         
         fire_output = image.copy()