Ritikumar/real_time_face_segmentation
0
1import os
2import cv2
3import numpy as np
4import streamlit as st
5import tensorflow as tf
6from PIL import Image
7
8
9st.set_page_config(page_title="Real-Time Face Segmentation", page_icon="š", layout="wide")
10
11
12IMG_SIZE = 256
13
14
15# =========================
16# FIND WEIGHTS AUTOMATICALLY
17# =========================
18def get_weights_path():
19 possible_paths = [
20 "unetmobilenetv2best.weights.h5",
21 "unet_mobilenetv2_best.weights.h5",
22 os.path.join("savedmodels", "unetmobilenetv2best.weights.h5"),
23 os.path.join("savedmodels", "unet_mobilenetv2_best.weights.h5"),
24 "/content/drive/MyDrive/Real-Time Face Segmentation for Movie Cast/savedmodels/unetmobilenetv2best.weights.h5",
25 "/content/drive/MyDrive/Real-Time Face Segmentation for Movie Cast/savedmodels/unet_mobilenetv2_best.weights.h5",
26 ]
27
28 for path in possible_paths:
29 if os.path.exists(path):
30 return path
31 return None
32
33
34def upsample(filters, size, apply_dropout=False):
35 initializer = tf.random_normal_initializer(0.0, 0.02)
36 result = tf.keras.Sequential()
37 result.add(
38 tf.keras.layers.Conv2DTranspose(
39 filters,
40 size,
41 strides=2,
42 padding="same",
43 kernel_initializer=initializer,
44 use_bias=False,
45 )
46 )
47 result.add(tf.keras.layers.BatchNormalization())
48 if apply_dropout:
49 result.add(tf.keras.layers.Dropout(0.3))
50 result.add(tf.keras.layers.ReLU())
51 return result
52
53
54def build_unet_mobilenetv2(input_shape=(256, 256, 3)):
55 base_model = tf.keras.applications.MobileNetV2(
56 input_shape=input_shape,
57 include_top=False,
58 weights="imagenet"
59 )
60
61 layer_names = [
62 "block_1_expand_relu",
63 "block_3_expand_relu",
64 "block_6_expand_relu",
65 "block_13_expand_relu",
66 "block_16_project"
67 ]
68
69 base_model_outputs = [base_model.get_layer(name).output for name in layer_names]
70 down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)
71 down_stack.trainable = False
72
73 up_stack = [
74 upsample(512, 3, apply_dropout=True),
75 upsample(256, 3, apply_dropout=True),
76 upsample(128, 3),
77 upsample(64, 3),
78 ]
79
80 inputs = tf.keras.layers.Input(shape=input_shape)
81 skips = down_stack(inputs)
82 x = skips[-1]
83 skips = reversed(skips[:-1])
84
85 for up, skip in zip(up_stack, skips):
86 x = up(x)
87 x = tf.keras.layers.Concatenate()([x, skip])
88
89 last = tf.keras.layers.Conv2DTranspose(
90 1, 3, strides=2, padding="same", activation="sigmoid"
91 )
92
93 outputs = last(x)
94 model = tf.keras.Model(inputs=inputs, outputs=outputs)
95 return model
96
97
98@st.cache_resource
99def load_segmentation_model():
100 weights_path = get_weights_path()
101 if weights_path is None:
102 raise FileNotFoundError(
103 "Weights file not found. Please place the file in app folder or savedmodels folder."
104 )
105
106 model = build_unet_mobilenetv2(input_shape=(IMG_SIZE, IMG_SIZE, 3))
107 model.load_weights(weights_path)
108 return model, weights_path
109
110
111def preprocess_image(uploaded_file):
112 image = Image.open(uploaded_file).convert("RGB")
113 original = np.array(image)
114 resized = cv2.resize(original, (IMG_SIZE, IMG_SIZE))
115 input_image = resized.astype(np.float32)
116 input_image = tf.keras.applications.mobilenet_v2.preprocess_input(input_image)
117 input_image = np.expand_dims(input_image, axis=0)
118 return original, resized, input_image
119
120
121def predict_mask(model, input_image, threshold=0.5):
122 pred = model.predict(input_image, verbose=0)[0]
123 pred_mask = pred.squeeze()
124 binary_mask = (pred_mask > threshold).astype(np.uint8) * 255
125 return pred_mask, binary_mask
126
127
128def resize_mask_to_original(mask, original_shape):
129 h, w = original_shape[:2]
130 return cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
131
132
133# ======== NEW: create bounding boxes from mask ========
134def create_box_image(original_image, mask):
135 """
136 mask: binary mask (0 / 255) same size as original_image
137 """
138 h, w = original_image.shape[:2]
139
140 # thoda clean mask
141 kernel = np.ones((3, 3), np.uint8)
142 mask_clean = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
143 mask_clean = cv2.dilate(mask_clean, kernel, iterations=1)
144
145 contours, _ = cv2.findContours(mask_clean, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
146
147 boxed = original_image.copy()
148
149 for cnt in contours:
150 area = cv2.contourArea(cnt)
151 if area < 100: # noise hatao
152 continue
153 x, y, w_box, h_box = cv2.boundingRect(cnt)
154 cv2.rectangle(boxed, (x, y), (x + w_box, y + h_box), (0, 255, 0), 2)
155
156 return boxed
157# ======================================================
158
159
160# =========================
161# SIDEBAR
162# =========================
163st.sidebar.title("Navigation")
164page = st.sidebar.radio(
165 "",
166 ["Introduction", "How it Works", "Prediction", "About"]
167)
168
169
170# =========================
171# INTRODUCTION PAGE
172# =========================
173if page == "Introduction":
174 st.header("š Real-Time Face Segmentation")
175 st.markdown(
176 '<div class="sub-text">This application segments the face region from an uploaded image and generates a binary mask and an overlay result using a U-Net model with a MobileNetV2 encoder.</div>',
177 unsafe_allow_html=True
178 )
179
180 st.markdown('<div class="section-card">', unsafe_allow_html=True)
181 st.markdown('<div class="small-title">Project Overview</div>', unsafe_allow_html=True)
182 st.write(
183 "This project focuses on face segmentation, where the goal is to identify the facial region at the pixel level. "
184 "Unlike simple image classification, segmentation predicts a mask that marks the exact area of interest in the image. "
185 "The model used here combines the U-Net decoder with a pretrained MobileNetV2 encoder for efficient and accurate prediction."
186 )
187 st.markdown('</div>', unsafe_allow_html=True)
188
189 st.markdown('<div class="section-card">', unsafe_allow_html=True)
190 st.markdown('<div class="small-title">What this app does</div>', unsafe_allow_html=True)
191 st.write(
192 "- Accepts an image upload\n"
193 "- Preprocesses the image to match model input format\n"
194 "- Predicts a face segmentation mask\n"
195 "- Displays the original image, binary mask, and bounding-box result"
196 )
197 st.markdown('</div>', unsafe_allow_html=True)
198
199 st.markdown('<div class="highlight-box">', unsafe_allow_html=True)
200 st.write(
201 "**Model note:** The app loads a weights-only checkpoint, so the exact same U-Net + MobileNetV2 architecture is recreated before loading the saved weights."
202 )
203 st.markdown('</div>', unsafe_allow_html=True)
204
205
206# =========================
207# HOW IT WORKS PAGE
208# =========================
209elif page == "How it Works":
210 st.title("āļø How it Works")
211
212 st.markdown("""
213 This project performs **face segmentation**, where the model predicts the face region at pixel level.
214 The input image is resized to **256x256**, passed through the trained **U-Net + MobileNetV2** model,
215 and the output is a predicted mask that highlights the facial area.
216 """)
217
218 st.subheader("Pipeline")
219 st.write("""
220 1. User uploads an image
221 2. Image is resized to 256x256
222 3. MobileNetV2 preprocessing is applied
223 4. Image goes into encoder layers
224 5. Decoder reconstructs segmentation mask
225 6. Threshold is applied to create binary mask
226 7. Mask is resized back to original image size
227 8. Bounding box is drawn on original image
228 """)
229
230 st.subheader("Architecture Flow")
231
232 col1, col2, col3, col4 = st.columns(4)
233
234 with col1:
235 st.markdown(
236 """
237 <div style="
238 background-color:#e3f2fd;
239 padding:20px;
240 border-radius:12px;
241 text-align:center;
242 font-weight:bold;
243 color:#0d47a1;
244 min-height:120px;
245 display:flex;
246 align-items:center;
247 justify-content:center;
248 ">
249 Input Image<br>(256 x 256 x 3)
250 </div>
251 """,
252 unsafe_allow_html=True
253 )
254
255 with col2:
256 st.markdown(
257 """
258 <div style="
259 background-color:#ede7f6;
260 padding:20px;
261 border-radius:12px;
262 text-align:center;
263 font-weight:bold;
264 color:#4a148c;
265 min-height:120px;
266 display:flex;
267 align-items:center;
268 justify-content:center;
269 ">
270 Encoder<br>MobileNetV2
271 </div>
272 """,
273 unsafe_allow_html=True
274 )
275
276 with col3:
277 st.markdown(
278 """
279 <div style="
280 background-color:#e8f5e9;
281 padding:20px;
282 border-radius:12px;
283 text-align:center;
284 font-weight:bold;
285 color:#1b5e20;
286 min-height:120px;
287 display:flex;
288 align-items:center;
289 justify-content:center;
290 ">
291 Decoder<br>U-Net Upsampling
292 </div>
293 """,
294 unsafe_allow_html=True
295 )
296
297 with col4:
298 st.markdown(
299 """
300 <div style="
301 background-color:#fff3e0;
302 padding:20px;
303 border-radius:12px;
304 text-align:center;
305 font-weight:bold;
306 color:#e65100;
307 min-height:120px;
308 display:flex;
309 align-items:center;
310 justify-content:center;
311 ">
312 Output Mask<br>(256 x 256 x 1)
313 </div>
314 """,
315 unsafe_allow_html=True
316 )
317
318 st.markdown("<br>", unsafe_allow_html=True)
319
320 st.markdown(
321 """
322 <div style="
323 background-color:#f5f5f5;
324 padding:18px;
325 border-radius:12px;
326 border-left:5px solid #616161;
327 ">
328 <b>Simple Explanation:</b><br><br>
329 - <b>Input Image:</b> Original face image is given to the model.<br>
330 - <b>Encoder:</b> MobileNetV2 extracts important visual features from the image.<br>
331 - <b>Decoder:</b> U-Net upsamples those features to rebuild the face region pixel by pixel.<br>
332 - <b>Output Mask:</b> Final segmentation mask highlights the predicted facial area.
333 </div>
334 """,
335 unsafe_allow_html=True
336 )
337
338 st.subheader("Why U-Net + MobileNetV2?")
339 st.write("""
340 - **MobileNetV2** is lightweight and good at extracting image features
341 - **U-Net** is powerful for segmentation because it restores spatial details
342 - Together, they provide a good balance of speed and segmentation quality
343 """)
344
345
346# =========================
347# PREDICTION PAGE
348# =========================
349elif page == "Prediction":
350 st.title("š¼ļø Prediction")
351
352 uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
353 threshold = st.slider("Mask Threshold", 0.1, 0.9, 0.5, 0.05)
354
355 if uploaded_file is not None:
356 try:
357 with st.spinner("Loading model and generating prediction..."):
358 model, used_weights_path = load_segmentation_model()
359 original_img, resized_img, input_img = preprocess_image(uploaded_file)
360 pred_mask, binary_mask = predict_mask(model, input_img, threshold=threshold)
361 binary_mask_original = resize_mask_to_original(binary_mask, original_img.shape)
362 boxed_img = create_box_image(original_img, binary_mask_original)
363
364 st.success(f"Prediction completed. Loaded weights from: {used_weights_path}")
365
366 c1, c2, c3 = st.columns(3)
367 with c1:
368 st.image(original_img, caption="Original Image", use_container_width=True)
369 with c2:
370 st.image(binary_mask_original, caption="Predicted Mask", use_container_width=True)
371 with c3:
372 st.image(boxed_img, caption="Bounding Box Result", use_container_width=True)
373
374 except Exception as e:
375 st.error(f"Error while loading model or generating prediction: {e}")
376
377
378# =========================
379# ABOUT PAGE
380# =========================
381elif page == "About":
382 st.markdown('<div class="main-title">ā¹ļø About</div>', unsafe_allow_html=True)
383
384 st.markdown('<div class="section-card">', unsafe_allow_html=True)
385 st.markdown('<div class="small-title">Project Description</div>', unsafe_allow_html=True)
386 st.write(
387 "This application was developed as a deep learning project for face segmentation. "
388 "Its purpose is to demonstrate how a convolutional neural network can identify the face region in an image and represent it as a segmentation mask."
389 )
390 st.markdown('</div>', unsafe_allow_html=True)
391
392 st.markdown('<div class="section-card">', unsafe_allow_html=True)
393 st.markdown('<div class="small-title">Technical Details</div>', unsafe_allow_html=True)
394 st.write(
395 "- Model: U-Net with MobileNetV2 encoder\n"
396 "- Input size: 256 x 256 x 3\n"
397 "- Output: 256 x 256 x 1 binary face mask\n"
398 "- Framework: TensorFlow / Keras\n"
399 "- App framework: Streamlit"
400 )
401 st.markdown('</div>', unsafe_allow_html=True)
402
403 st.markdown('<div class="section-card">', unsafe_allow_html=True)
404 st.markdown('<div class="small-title">Purpose</div>', unsafe_allow_html=True)
405 st.write(
406 "The goal of this project is to show the complete deployment flow of a segmentation model, "
407 "from preprocessing and model loading to prediction and visualization in a simple user interface."
408 )
409 st.markdown('</div>', unsafe_allow_html=True)
410
411
412st.markdown("<br><div class='footer-note'>Built for face segmentation using deep learning.</div>", unsafe_allow_html=True)