Update app.py

app.py

@@ -29,8 +29,6 @@ def visualize_model_output(prediction, img):
 
     output = np.zeros(prediction.shape)
 
-    print(output.shape,'shape of output')
-
     for unq_class in unique_classes:
         print(unq_class,'unq_class')
         rgb_class_unique = rgb_colors[str(int(unq_class))]
@@ -45,13 +43,9 @@ def visualize_model_output(prediction, img):
     output = output.astype(np.int32)
     img = img.astype(np.int32)
 
-    print(output,'shapes')
-    print(img,'shapes2')
-
     
     
     added_image = cv2.addWeighted(img,0.5,output,0.1,0)
-    print(added_image)
     return added_image
 
 
@@ -62,8 +56,18 @@ def do_prediction(model_name, img):
     match model_name:
         # numerical output
         case "SBB/eynollah-column-classifier": 
+
+            img_1ch = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+
+            img_1ch = img_1ch / 255.0
+            img_1ch = cv2.resize(img_1ch, (448, 448), interpolation=cv2.INTER_NEAREST)
+            img_in = np.zeros((1, img_1ch.shape[0], img_1ch.shape[1], 3))
+            img_in[0, :, :, 0] = img_1ch[:, :]
+            img_in[0, :, :, 1] = img_1ch[:, :]
+            img_in[0, :, :, 2] = img_1ch[:, :]
                       
-            num_col=model.layers[len(model.layers)-1].output_shape[1]
+            label_p_pred = model.predict(img_in, verbose=0)
+            num_col = np.argmax(label_p_pred[0]) + 1
             return "Found {} columns".format(num_col), None
                             
         # bitmap output