Update app.py

app.py

@@ -1,6 +1,5 @@
 import os
-# Force Keras 2 logic to prevent the 'recursion depth' and 'quantization' errors 
-# common when loading Kaggle-trained .h5 files in new environments.
+# Force Keras 2 logic to prevent recursion/quantization errors from Kaggle .h5 files
 os.environ["TF_USE_LEGACY_KERAS"] = "1" 
 
 import gradio as gr
@@ -15,79 +14,85 @@ from huggingface_hub import hf_hub_download
 REPO_ID = "mediaportal/Roadsegmentation" 
 MODEL_FILENAME = "trained_model_33_cpu.h5"
 
-# If your repo is PRIVATE, add your token to Space Secrets as 'HF_TOKEN'
+# BDD100K Color Dictionary from your notebook
+COLOR_DICT = {
+    0: (128, 128, 128),  # road - gray
+    1: (230, 230, 50),   # sidewalk - yellow
+    8: (50, 150, 50),    # vegetation - green
+    10: (128, 180, 255), # sky - blue
+    11: (255, 0, 0),     # person - red
+    13: (0, 0, 255),     # car - blue
+    19: (0, 0, 0)        # unknown - black
+}
+
 hf_token = os.getenv("HF_TOKEN")
 model = None
 
 def load_model():
     global model
     try:
-        # Download the model file
         path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME, token=hf_token)
-        
-        # Load using the Classic Keras engine
-        # compile=False is required because segmentation models often use 
-        # custom Loss functions (like IoU or Dice) that are hard to reload.
+        # compile=False is used because the notebook uses SparseCategoricalCrossentropy
         model = keras.models.load_model(path, compile=False)
         return "✅ Road Segmentation Model Loaded"
     except Exception as e:
         return f"❌ Error: {str(e)}"
 
-def predict_segmentation(img):
+def segment_road(img):
     if model is None:
-        return None
+        return None, None
     
-    # 1. Store original dimensions
-    h, w = img.shape[:2]
+    # 1. Store original size for scaling back
+    h_orig, w_orig = img.shape[:2]
     
-    # 2. Preprocessing
-    # BDD100K segmentation models typically use 256x256 or 512x512.
-    # We resize to 256x256 based on common CPU-optimized configurations.
-    input_size = (256, 256)
-    img_resized = cv2.resize(img, input_size)
+    # 2. Preprocessing (Notebook uses 192 height, 256 width)
+    img_resized = cv2.resize(img, (256, 192)) 
     img_array = img_resized.astype('float32') / 255.0
     img_array = np.expand_dims(img_array, axis=0)
     
-    # 3. Predict the Mask
+    # 3. Predict (Returns 20 channels for 20 classes)
     prediction = model.predict(img_array)[0]
     
-    # 4. Process the Mask
-    # The output is a probability map. Threshold at 0.5 to get binary road/not-road.
-    mask = (prediction > 0.5).astype(np.uint8) * 255
-    if len(mask.shape) == 3:
-        mask = np.squeeze(mask, axis=-1)
-        
-    # Resize mask back to original image size
-    mask_full = cv2.resize(mask, (w, h), interpolation=cv2.INTER_NEAREST)
+    # Get the class index with the highest probability for each pixel
+    mask = np.argmax(prediction, axis=-1).astype(np.uint8)
+    
+    # 4. Create Outputs
+    # A. Full Semantic Map (Colorizing all classes)
+    full_mask_color = np.zeros((192, 256, 3), dtype=np.uint8)
+    for class_idx, color in COLOR_DICT.items():
+        full_mask_color[mask == class_idx] = color
+    
+    # B. Road Highlight Overlay (Class 0 is Road)
+    road_mask = (mask == 0).astype(np.uint8) * 255
+    road_mask_resized = cv2.resize(road_mask, (w_orig, h_orig), interpolation=cv2.INTER_NEAREST)
     
-    # 5. Create the Green Overlay
-    # We create a green version of the original image
     overlay = img.copy()
-    overlay[mask_full > 0] = [0, 255, 0] # Apply green color (RGB)
+    overlay[road_mask_resized > 0] = [0, 255, 0] # Highlight road in green
     
-    # Blend original and green overlay (0.6 original + 0.4 green)
-    output = cv2.addWeighted(img, 0.6, overlay, 0.4, 0)
+    # Blend: 70% original image, 30% green highlight
+    highlighted_road = cv2.addWeighted(img, 0.7, overlay, 0.3, 0)
     
-    return output
+    # Resize full mask back to original aspect ratio for display
+    full_mask_resized = cv2.resize(full_mask_color, (w_orig, h_orig), interpolation=cv2.INTER_NEAREST)
+    
+    return highlighted_road, full_mask_resized
 
 # --- GRADIO INTERFACE ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚗 ADAS: Road Surface Segmentation")
-    gr.Markdown("Upload a dashboard camera image to visualize the drivable road area detection.")
+    gr.Markdown("# 🚗 ADAS Road & Scene Segmentation")
+    gr.Markdown("Upload a dashboard image to identify the drivable road surface and other objects.")
     
     status = gr.Markdown("⏳ Initializing system...")
     
     with gr.Row():
-        input_img = gr.Image(label="Original View", type="numpy")
-        output_img = gr.Image(label="Detected Road (Green Overlay)")
+        input_img = gr.Image(label="Input Dashboard Image", type="numpy")
+        output_overlay = gr.Image(label="Drivable Road (Green Highlight)")
+        output_full = gr.Image(label="Full Semantic Map")
     
-    btn = gr.Button("Analyze Road", variant="primary")
+    btn = gr.Button("Analyze Scene", variant="primary")
     
-    # Automatically load model on start
     demo.load(load_model, outputs=status)
-    
-    # Connect button to prediction
-    btn.click(fn=predict_segmentation, inputs=input_img, outputs=output_img)
+    btn.click(fn=segment_road, inputs=input_img, outputs=[output_overlay, output_full])
 
 if __name__ == "__main__":
     demo.queue().launch()