changes

scripts/welspun_classifier.py

@@ -43,7 +43,7 @@ class Welspun_Classifier(ModelWrapper):
 
         # open-vino
         self.sink_layer = {'0': 'conv2d_58/BiasAdd/Add', '1': 'conv2d_66/BiasAdd/Add', '2': 'conv2d_74/BiasAdd/Add'}
-        self.model_detector_pth = os.path.join(self.base_model_path, "resnet34-wel4.xml")
+        self.model_detector_pth = os.path.join(self.base_model_path, "resnet34-demo.xml")
         self.model_bin = os.path.splitext(self.model_detector_pth)[0] + ".bin"
         self.ie = IECore()
         self.net1 = self.ie.read_network(model=self.model_detector_pth, weights=self.model_bin)
@@ -222,12 +222,13 @@ class Welspun_Classifier(ModelWrapper):
         x.append(prob[1])
         x.append(prob[2])
         x.append(prob[3])
+        x.append(prob[4])
         a = x.index(max(x))
         # print(type(prob))
         # if self.counter%1 == 0:
         self.counter = self.counter + 1
 
-        if a == 0 and x[0] > 0.95:
+        if a == 1 and x[1] > 0.95:
             if self.defect_type == 'Mix':
                 cv2.putText(frame, text="Mix Color Defect Detected", org=(50, 50),
                             color=(0, 0, 255),
@@ -251,7 +252,7 @@ class Welspun_Classifier(ModelWrapper):
                 #                       "sound_1")
                 #     logger.info(f"Probability: {prob}")
                 #     self.counter = 0
-        elif a == 2 and x[2] > 0.95:
+        elif a == 3 and x[3] > 0.95:
             if self.defect_type == 'Short':
                 cv2.putText(frame, text="Short Defect Detected", org=(50, 50),
                             color=(0, 0, 255),
@@ -275,7 +276,7 @@ class Welspun_Classifier(ModelWrapper):
             #                       "sound_1")
             #     logger.info(f"Probability: {prob}")
             #     self.counter = 0
-        elif a == 3 and x[3] > 0.95:
+        elif a == 4 and x[4] > 0.95:
             if self.defect_type == 'Split':
                 cv2.putText(frame, text="Split Defect Detected", org=(50, 50),
                             color=(0, 0, 255),
@@ -297,6 +298,28 @@ class Welspun_Classifier(ModelWrapper):
             #                       "sound_1")
             #     logger.info(f"Probability: {prob}")
             #     self.counter = 0
+        elif a == 0 and x[0] > 0.95:
+            if self.defect_type == 'good':
+                cv2.putText(frame, text="No Defect", org=(50, 50),
+                            color=(0, 0, 255),
+                            thickness=2,
+                            fontScale=1, fontFace=cv2.LINE_AA)
+                # pass
+            else:
+                self.defect_type = 'good'
+                resized_frame = cv2.resize(frame, (64, 64))
+                cv2.putText(frame, text="No Defect", org=(50, 50),
+                            color=(0, 0, 255),
+                            thickness=2,
+                            fontScale=1, fontFace=cv2.LINE_AA)
+                self.send_payload("No Defect", resized_frame, "no defect " + str(prob[3]), "#472020", "#ed2020",
+                                  "sound_1")
+                logger.info(f"Probability: {prob}")
+            # if self.counter % 25 == 0:
+            #     self.send_payload("Split Defect Detected", resized_frame, "Split " + str(prob[3]), "#472020", "#ed2020",
+            #                       "sound_1")
+            #     logger.info(f"Probability: {prob}")
+            #     self.counter = 0
         elif a == 2:
             pass