Update app.py

app.py

@@ -9,12 +9,32 @@ from huggingface_hub import from_pretrained_keras
 model = from_pretrained_keras("Harveenchadha/low-light-image-enhancement", compile=False)
 #examples = ['examples/179.png', 'examples/493.png', 'examples/780.png']
 
-
+def get_enhanced_image(data, output):
+    r1 = output[:, :, :, :3]
+    r2 = output[:, :, :, 3:6]
+    r3 = output[:, :, :, 6:9]
+    r4 = output[:, :, :, 9:12]
+    r5 = output[:, :, :, 12:15]
+    r6 = output[:, :, :, 15:18]
+    r7 = output[:, :, :, 18:21]
+    r8 = output[:, :, :, 21:24]
+    x = data + r1 * (tf.square(data) - data)
+    x = x + r2 * (tf.square(x) - x)
+    x = x + r3 * (tf.square(x) - x)
+    enhanced_image = x + r4 * (tf.square(x) - x)
+    x = enhanced_image + r5 * (tf.square(enhanced_image) - enhanced_image)
+    x = x + r6 * (tf.square(x) - x)
+    x = x + r7 * (tf.square(x) - x)
+    enhanced_image = x + r8 * (tf.square(x) - x)
+    return enhanced_image
+    
+    
 def infer(original_image):
     image = keras.preprocessing.image.img_to_array(original_image)
     image = image.astype("float32") / 255.0
     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)