Create app.py

app.py

@@ -0,0 +1,181 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from transformers import SegformerForSemanticSegmentation
+from torchvision import transforms
+from PIL import Image
+import numpy as np
+import os
+
+# 1. SETUP & MODEL LOADING
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+model = SegformerForSemanticSegmentation.from_pretrained(
+    "nvidia/mit-b2", 
+    num_labels=4, 
+    id2label={0: "Soil", 1: "Bedrock", 2: "Sand", 3: "Big Rock"},
+    label2id={"Soil": 0, "Bedrock": 1, "Sand": 2, "Big Rock": 3},
+    ignore_mismatched_sizes=True
+)
+
+# Load your Stirling weights (Must be in the same folder as app.py)
+try:
+    checkpoint = torch.load('SegFormer_B2_Final_Stirling_3456526.pth', map_location=device)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    print("Model Weights Loaded Successfully")
+except Exception as e:
+    print(f"⚠️ Weights missing or error: {e}. Running with base weights.")
+
+model.to(device).eval()
+
+# 2. COLOR MAP (Brightened for visibility)
+COLOR_MAP = {
+    0: [0, 255, 0],      # Soil (Neon Green) - UPDATED
+    1: [0, 0, 255],      # Bedrock (Electric Blue) - UPDATED
+    2: [255, 215, 0],    # Sand (Yellow)
+    3: [255, 0, 0],      # Big Rock (Red)
+    -1: [0, 0, 0]
+}
+
+def apply_mask_safe(preds, folder, img_path, suffix, w, h):
+    """
+    Finds the mask by searching for the 9-digit SCLK ID inside the folder,
+    ignoring 'EDR', 'NLA', or other naming variations.
+    """
+    filename = os.path.basename(img_path)
+    
+    # 1. Extract the 9-digit numeric ID (SCLK)
+    # Example: NLA_601686301EDR... -> 601686301
+    import re
+    match = re.search(r'\d{9}', filename)
+    if not match:
+        print(f"❌ Could not find a 9-digit ID in {filename}")
+        return preds
+    
+    seq_id = match.group(0)
+    print(f"DEBUG: Searching for ID {seq_id} in {folder}...")
+
+    # 2. Search the folder for a file containing this ID and the suffix (mxy/rng)
+    target_file = None
+    if os.path.exists(folder):
+        for f in os.listdir(folder):
+            # Check if the 9-digit ID is in the filename AND it's a .png
+            if seq_id in f and f.lower().endswith('.png'):
+                # Also check if it matches the specific suffix if needed
+                if suffix in f.lower():
+                    target_file = f
+                    break
+    
+    # 3. Apply if found
+    if target_file:
+        path = os.path.join(folder, target_file)
+        mask = Image.open(path).convert('L').resize((w, h), Image.NEAREST)
+        mask_np = np.array(mask)
+        preds[mask_np > 0] = -1
+        print(f"✅ SUCCESS: Applied mask from {target_file}")
+    else:
+        print(f"❌ FAIL: No mask for ID {seq_id} found in {folder}")
+        
+    return preds
+
+def segment_mars(img_path):
+    if not img_path: return None
+    
+    # 1. Load Image
+    raw_img = Image.open(img_path).convert('RGB')
+    orig_w, orig_h = raw_img.size
+    
+    # 2. Inference (SegFormer)
+    preprocess = transforms.Compose([
+        transforms.Resize((256, 256)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+    ])
+    input_tensor = preprocess(raw_img).unsqueeze(0).to(device)
+    
+    with torch.no_grad():
+        outputs = model(pixel_values=input_tensor)
+        logits = F.interpolate(outputs.logits, size=(orig_h, orig_w), mode='bilinear')
+        probs = F.softmax(logits, dim=1)
+        confidences, preds = torch.max(probs, dim=1)
+        preds = preds.squeeze().cpu().numpy()
+
+    # 3. DIRECT FILENAME SWAP (EDR -> MXY / EDR -> RNG)
+    filename = os.path.basename(img_path)
+    
+    # Generate mask names by replacing 'EDR' with 'MXY' or 'RNG'
+    # and changing extension to .png
+    mxy_filename = filename.replace("EDR", "MXY").replace(".JPG", ".png").replace(".jpg", ".png")
+    rng_filename = filename.replace("EDR", "RNG").replace(".JPG", ".png").replace(".jpg", ".png")
+
+    mxy_path = os.path.join("stirling_masks_bundle", "rover_mxy", mxy_filename)
+    rng_path = os.path.join("stirling_masks_bundle", "range_rng", rng_filename)
+
+    # Apply MXY Mask
+    if os.path.exists(mxy_path):
+        mxy = Image.open(mxy_path).convert('L').resize((orig_w, orig_h), Image.NEAREST)
+        preds[np.array(mxy) > 0] = -1
+        
+    else:
+        print(f"❌ MXY Not Found: {mxy_path}")
+
+    # Apply RNG Mask
+    if os.path.exists(rng_path):
+        rng = Image.open(rng_path).convert('L').resize((orig_w, orig_h), Image.NEAREST)
+        preds[np.array(rng) > 0] = -1
+    else:
+        print(f"❌ RNG Not Found: {rng_path}")
+
+    # 4. OVERLAY
+    mask_rgb = np.zeros((orig_h, orig_w, 3), dtype=np.uint8)
+    for cls_id, color in COLOR_MAP.items():
+        mask_rgb[preds == cls_id] = color
+        
+    overlay = (np.array(raw_img) * 0.5 + mask_rgb * 0.5).astype(np.uint8)
+    return Image.fromarray(overlay)
+
+# 3. CUSTOM HTML LEGEND
+legend_html = """
+<div style="display: flex; justify-content: center; gap: 20px; font-weight: bold; margin-bottom: 10px;">
+    <div style="display: flex; align-items: center; gap: 5px;">
+        <div style="width: 20px; height: 20px; background-color: rgb(0, 255, 0); border: 1px solid white;"></div> <span>Soil</span>
+    </div>
+    <div style="display: flex; align-items: center; gap: 5px;">
+        <div style="width: 20px; height: 20px; background-color: rgb(0, 0, 255); border: 1px solid white;"></div> <span>Bedrock</span>
+    </div>
+    <div style="display: flex; align-items: center; gap: 5px;">
+        <div style="width: 20px; height: 20px; background-color: rgb(255, 215, 0); border: 1px solid white;"></div> <span>Sand</span>
+    </div>
+    <div style="display: flex; align-items: center; gap: 5px;">
+        <div style="width: 20px; height: 20px; background-color: rgb(255, 0, 0); border: 1px solid white;"></div> <span>Big Rock</span>
+    </div>
+    <div style="display: flex; align-items: center; gap: 5px;">
+        <div style="width: 20px; height: 20px; background-color: rgb(0, 0, 0); border: 1px solid white;"></div> <span>Rover / Background Mask</span>
+    </div>
+</div>
+"""
+# 3. GRADIO INTERFACE
+with gr.Blocks() as demo:
+    gr.Markdown(f"## NASA AI4Mars Expert Fused Classifier")
+    gr.HTML(legend_html)
+    with gr.Row():
+        img_input = gr.Image(type="filepath", label="Input Martian Image",interactive=False)
+        img_output = gr.Image(type="pil", label="Fused Ground Truth Prediction")
+
+    btn = gr.Button("Execute Data Fusion Segmentation")
+    btn.click(segment_mars, inputs=img_input, outputs=img_output)
+
+    gr.Examples(
+        examples=[
+            ["NLA_601686301EDR_F0732112NCAM00353M1.JPG"],
+            ["NLB_436292094EDR_F0211028NCAM00257M1.JPG"],
+            ["NLB_486005519EDR_F0481570NCAM07813M1.JPG"],
+            ["NLB_519658137EDR_F0550000NCAM00654M1.JPG"],
+            ["NLB_541230242EDR_F0611140NCAM07753M1.JPG"],
+            ["NLB_621571338EDR_F0763002NCAM00207M1.JPG"]
+        ],
+        inputs=img_input
+    )
+
+if __name__ == "__main__":
+    demo.launch(share=True)