add documentation in app.py

app.py

@@ -1,3 +1,11 @@
+'''
+This Python script is a web application that performs human body part segmentation 
+using a pre-trained deep learning model called DeepLabv3+. 
+The application is built using the Streamlit library and uses the Hugging Face Hub 
+to download the pre-trained model.
+'''
+
+# import libraries
 import numpy as np
 import tensorflow as tf
 import streamlit as st
@@ -5,25 +13,34 @@ from PIL import Image
 from huggingface_hub import from_pretrained_keras
 import cv2
 
-
+# The model used is the DeepLabv3+ model with a ResNet50 backbone.
 model = from_pretrained_keras("keras-io/deeplabv3p-resnet50")
 
+# A colormap is defined to map the predicted segmentation masks to colors for better visualization
 colormap = np.array([[0,0,0], [31,119,180], [44,160,44], [44, 127, 125], [52, 225, 143],
                     [217, 222, 163], [254, 128, 37], [130, 162, 128], [121, 7, 166], [136, 183, 248],
                     [85, 1, 76], [22, 23, 62], [159, 50, 15], [101, 93, 152], [252, 229, 92],
                     [167, 173, 17], [218, 252, 252], [238, 126, 197], [116, 157, 140], [214, 220, 252]], dtype=np.uint8)
-                    
+
+# size of the input image is defined as 512x512 pixels                    
 img_size = 512
                     
 def read_image(image):
+    '''
+    read_image: reads in the input image and preprocesses it 
+                by resizing it to the defined size and normalizing it to values between -1 and 1
+    '''
     image = tf.convert_to_tensor(image)
     image.set_shape([None, None, 3])
     image = tf.image.resize(images=image, size=[img_size, img_size])
-    image = image / 127.5 - 1
+    image = image / 255
     return image
 
 
 def infer(model, image_tensor):
+    '''
+    infer: performs inference using the pre-trained model and returns the predicted segmentation mask.
+    '''
     predictions = model.predict(np.expand_dims((image_tensor), axis=0))
     predictions = np.squeeze(predictions)
     predictions = np.argmax(predictions, axis=2)
@@ -31,6 +48,10 @@ def infer(model, image_tensor):
 
 
 def decode_segmentation_masks(mask, colormap, n_classes):
+    '''
+     decode_segmentation_masks: maps the predicted segmentation mask to the defined colormap 
+                                to produce a colored mask.
+     '''
     r = np.zeros_like(mask).astype(np.uint8)
     g = np.zeros_like(mask).astype(np.uint8)
     b = np.zeros_like(mask).astype(np.uint8)
@@ -44,6 +65,9 @@ def decode_segmentation_masks(mask, colormap, n_classes):
 
 
 def get_overlay(image, colored_mask):
+    '''
+    get_overlay: overlays the colored mask on the original image for visualization
+    '''
     image = tf.keras.preprocessing.image.array_to_img(image)
     image = np.array(image).astype(np.uint8)
     overlay = cv2.addWeighted(image, 0.35, colored_mask, 0.65, 0)
@@ -51,6 +75,14 @@ def get_overlay(image, colored_mask):
 
 
 def segmentation(input_image):
+    '''
+    segmentation: 
+        returns,
+           - prediction_colormap: function is used to convert the prediction mask into a colored mask, 
+             where each class is assigned a unique color from a predefined color map.
+
+           - overlay: used to create an overlay image by blending the original input image with the colored mask
+    '''
     image_tensor = read_image(input_image)
     prediction_mask = infer(image_tensor=image_tensor, model=model)
     prediction_colormap = decode_segmentation_masks(prediction_mask, colormap, 20)
@@ -58,7 +90,10 @@ def segmentation(input_image):
     return (overlay, prediction_colormap)
 
 
+## Streamlit interface
+
 st.header("Segmentación de partes del cuerpo humano")
+
 st.subheader("Demo de Spaces usando Streamlit y segmentacion de imagenes [Space original](https://huggingface.co/spaces/PKaushik/Human-Part-Segmentation)")
 
 st.markdown("Sube una imagen o selecciona un ejemplo para segmentar las distintas partes del cuerpo humano")
@@ -102,4 +137,7 @@ if file_imagen is not None:
         with col2:
             st.subheader("Mask: ")
             st.image(output[1], width=425)
-            
+else:
+    st.write("Ocurrió un error al segmentar la imagen, intenta con otra imagen")
+
+