Update app.py

app.py

@@ -481,7 +481,15 @@ plt.show()
 #plt.imsave(name,pre)
 #cv2.imshow(pre)
 
-def result(Input,Choice,Step=300):
+def ssim(original,predict):
+   ssim = tf.image.ssim(original, predict, max_val=1.0, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03)
+   return ssim
+
+def psnr(Input,Output,Choice):
+  psnr = tf.image.psnr(Input, Output, max_val=255)
+  return psnr
+  
+def result(Input,Choice):
     if Choice=="Coloring":
       pre_trained = tf.keras.models.load_model("gradio_pix2pix.h5")
       size0 = Input.shape[0]
@@ -490,11 +498,12 @@ def result(Input,Choice,Step=300):
       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
+      Output = prediction[0]
+      Output = (Output+1)*127.5
+      Output = np.uint8(Output)
+      Output = cv2.resize(Output, (size1,size0), interpolation = cv2.INTER_AREA)
+      psnr = psnr(Input,Output,"Coloring")
+      return Output,psnr
       
     if Choice=="Enhancement":
       pre_trained2 = tf.keras.models.load_model("gradio_pix2pix.h5")
@@ -510,15 +519,6 @@ def result(Input,Choice,Step=300):
       Input = cv2.resize(Input, (size1,size0), interpolation = cv2.INTER_AREA)
       return Input
 
-def ssim(original,predict):
-   ssim = tf.image.ssim(original, predict, max_val=1.0, filter_size=11, filter_sigma=1.5, k1=0.01, k2=0.03)
-   return ssim
-
-def psnr(Input,Choice):
-  Output = result(Input,Choice)
-  psnr = tf.image.psnr(Input, Output, max_val=255)
-  print(psnr)
-  return psnr
 
 #lst = cv2.imread('/content/drive/MyDrive/ColabNotebooks/enhance/low-sat.jpg')
 #r = result(lst)