Update app.py

app.py

@@ -447,18 +447,21 @@ pre_trained = tf.keras.models.load_model("gradio_pix2pix.h5")
 #plt.imsave(name,pre)
 #cv2.imshow(pre)
 
-def result(Input):
-    size0 = Input.shape[0]
-    size1 = Input.shape[1]
-    Input = cv2.resize(Input, (256,256), interpolation = cv2.INTER_AREA)
-    Input = cv2.cvtColor(Input , cv2.COLOR_BGR2GRAY)
-    Input = np.array(Input).reshape(1,256,256,1)
-    prediction = pre_trained(Input,training=True)
-    Input = prediction[0]
-    Input = (Input+1)*127.5
-    Input = np.uint8(Input)
-    Input = cv2.resize(Input, (size1,size0), interpolation = cv2.INTER_AREA)
-    return Input
+def result(Input,Choice):
+    if Choice=="Coloring":
+      size0 = Input.shape[0]
+      size1 = Input.shape[1]
+      Input = cv2.resize(Input, (256,256), interpolation = cv2.INTER_AREA)
+      Input = cv2.cvtColor(Input , cv2.COLOR_BGR2GRAY)
+      Input = np.array(Input).reshape(1,256,256,1)
+      prediction = pre_trained(Input,training=True)
+      Input = prediction[0]
+      Input = (Input+1)*127.5
+      Input = np.uint8(Input)
+      Input = cv2.resize(Input, (size1,size0), interpolation = cv2.INTER_AREA)
+      return Input
+    if Choice=="Enhancement":
+      return Input
 
 #lst = cv2.imread('/content/drive/MyDrive/ColabNotebooks/enhance/low-sat.jpg')
 #r = result(lst)