Update app.py

app.py

@@ -1,21 +1,42 @@
 import gradio as gr
 from PIL import Image
+import cv2
+import numpy as np
+from tensorflow.keras.models import load_model
+from tensorflow.keras.layers import Input, Conv2D
+from tensorflow.keras.models import Model
 
 size = 128
 
 def build_model(input_shape):
-    model = build_model(input_shape=(size, size, 1))
+    input_layer = Input(input_shape)
+    
+    s1, p1 = encoder_block(input_layer, 64)
+    s2, p2 = encoder_block(p1, 128)
+    s3, p3 = encoder_block(p2, 256)
+    s4, p4 = encoder_block(p3, 512)
+
+    b1 = conv_block(p4, 1024)
+
+    d1 = decoder_block(b1, s4, 512)
+    d2 = decoder_block(d1, s3, 256)
+    d3 = decoder_block(d2, s2, 128)
+    d4 = decoder_block(d3, s1, 64)
+    
+    output_layer = Conv2D(1, 1, padding="same", activation="sigmoid")(d4)
+    
+    model = Model(input_layer, output_layer, name="U-Net")
     return model
 
 model = build_model(input_shape=(size, size, 1))
 model.load_weights('BreastCancerSegmentation.h5')
 
-def preprocess_image(image, size: int=128):
-    image = cv2.resize(image, (size,size))
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)  
-    image = image/255.
+def preprocess_image(image, size=128):
+    image = cv2.resize(image, (size, size))
+    image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    image = image / 255.
     return image
-  
+
 def segment(image):
     image = preprocess_image(image, size=size)
     image = np.expand_dims(image, 0)
@@ -28,11 +49,11 @@ def segment(image):
     return mask_image
 
 if __name__ == "__main__":
-  gr.Interface(
-      fn=segment,
-      inputs="image",
-      outputs=gr.Image(type="pil", label="Breast Cancer Mask"),
-      examples = [["/content/benign(10).png"], ["/content/benign(109).png"]],
-      title = "Breast Cancer Ultrasound Image Segmentation",
-      description = "Check out this exciting development in the field of breast cancer diagnosis and treatment! A demo of Breast Cancer Ultrasound Image Segmentation has been developed. Upload image file, or try out one of the examples below!"
-      ).launch(share=True, debug=True)
+    gr.Interface(
+        fn=segment,
+        inputs="image",
+        outputs=gr.Image(type="pil", label="Breast Cancer Mask"),
+        examples=[["/content/benign(10).png"], ["/content/benign(109).png"]],
+        title="Breast Cancer Ultrasound Image Segmentation",
+        description="Check out this exciting development in the field of breast cancer diagnosis and treatment! A demo of Breast Cancer Ultrasound Image Segmentation has been developed. Upload an image file, or try out one of the examples below!"
+    ).launch(share=True, debug=True)