Update app.py

app.py

@@ -35,22 +35,8 @@ def infer(original_image):
     image = np.expand_dims(image, axis=0)
     output = model.predict(image)
     output = get_enhanced_image(image, output)
-    output_image = output[0] * 255.0
-    output_image = output_image.clip(0, 255)
-    print(output_image)
-    print(len(output_image))
-    print([len(a) for a in output_image])
-    output_image = output_image.reshape(
-        (np.shape(output_image)[0], np.shape(output_image)[1], 3)
-    )
-    output_image = np.uint32(output_image)
-    
-    return output_image
-    
-    # output_image = tf.cast((output[0, :, :, :] * 255), dtype=np.uint8)
-    
-    
-    # #output_image = Image.fromarray(output_image.numpy())
+    output_image = tf.cast((output[0, :, :, :] * 255), dtype=np.uint8)    
+    output_image = Image.fromarray(output_image.numpy())
     # output_image = output_image.numpy()
     # print(output_image.shape())
     return output_image