Update app.py

app.py

@@ -44,21 +44,19 @@ model = smp.DeepLabV3Plus(
 )
 
 checkpoint = torch.load(MODEL_PATH, map_location=DEVICE)
-
-if "model_state_dict" in checkpoint:
-    model.load_state_dict(checkpoint["model_state_dict"])
-else:
-    model.load_state_dict(checkpoint)
+model.load_state_dict(
+    checkpoint["model_state_dict"] if "model_state_dict" in checkpoint else checkpoint
+)
 
 model.to(DEVICE)
 model.eval()
 
-# ================================
-# HELPERS
-# ================================
 normalize = Normalize()
 to_tensor = ToTensorV2()
 
+# ================================
+# HELPERS
+# ================================
 def pad_to_16(img):
     h, w = img.shape[:2]
     new_h = (h + 15) // 16 * 16
@@ -76,11 +74,13 @@ def colorize_mask(mask):
     return color_mask
 
 # ================================
-# INFERENCE FUNCTION
+# INFERENCE
 # ================================
 def segment_image(image):
-    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    if image is None:
+        return None, ""
 
+    # Gradio provides RGB already
     padded, orig_h, orig_w = pad_to_16(image)
 
     img = normalize(image=padded)["image"]
@@ -96,6 +96,7 @@ def segment_image(image):
     color_mask = colorize_mask(pred_mask)
     overlay = cv2.addWeighted(image, 0.6, color_mask, 0.4, 0)
 
+    # Image-level classification
     vals, counts = np.unique(pred_mask, return_counts=True)
     vals = vals[vals > 0]
 
@@ -112,17 +113,17 @@ def segment_image(image):
 # ================================
 with gr.Blocks() as demo:
     gr.Markdown("# 🏗 Structural Defect Segmentation")
-    
+
     with gr.Tab("Image Upload"):
         input_img = gr.Image(type="numpy")
         output_img = gr.Image()
         output_text = gr.Textbox()
         btn = gr.Button("Run Segmentation")
         btn.click(segment_image, inputs=input_img, outputs=[output_img, output_text])
-    
-    with gr.Tab("Live Camera"):
-        cam = gr.Image(source="webcam", streaming=True)
+
+    with gr.Tab("Live Camera (Real-Time)"):
+        cam = gr.Image(sources=["webcam"], streaming=True, type="numpy")
         cam_out = gr.Image()
-        cam.stream(segment_image, inputs=cam, outputs=[cam_out])
+        cam.stream(lambda x: segment_image(x)[0], inputs=cam, outputs=cam_out)
 
 demo.launch()