Create app.py

app.py

@@ -0,0 +1,295 @@
+import os
+import cv2
+import numpy as np
+import torch
+import gradio as gr
+import segmentation_models_pytorch as smp
+from PIL import Image
+from glob import glob
+from pipeline.ImgOutlier import detect_outliers
+from pipeline.normalization import align_images
+
+# 检测是否在Hugging Face环境中运行
+HF_SPACE = os.environ.get('SPACE_ID') is not None
+
+# 尝试导入上传模块，如果不在HF环境中才需要
+if not HF_SPACE:
+    try:
+        from uploader.do_spaces import upload_mask
+    except ImportError:
+        def upload_mask(image, prefix=""):
+            return "上传模块未加载"
+
+# Global Configuration
+MODEL_PATHS = {
+    "Metal Marcy": "models/MM_best_model.pth",
+    "Silhouette Jaenette": "models/SJ_best_model.pth"
+}
+
+REFERENCE_VECTOR_PATHS = {
+    "Metal Marcy": "models/MM_mean.npy",
+    "Silhouette Jaenette": "models/SJ_mean.npy"
+}
+
+REFERENCE_IMAGE_DIRS = {
+    "Metal Marcy": "reference_images/MM",
+    "Silhouette Jaenette": "reference_images/SJ"
+}
+
+# Category names and color mapping
+CLASSES = ['background', 'cobbles', 'drysand', 'plant', 'sky', 'water', 'wetsand']
+COLORS = [
+    [0, 0, 0],        # background - black
+    [139, 137, 137],  # cobbles - dark gray
+    [255, 228, 181],  # drysand - light yellow
+    [0, 128, 0],      # plant - green
+    [135, 206, 235],  # sky - sky blue
+    [0, 0, 255],      # water - blue
+    [194, 178, 128]   # wetsand - sand brown
+]
+
+# Load model function
+def load_model(model_path, device="cuda"):
+    try:
+        # 如果在HF环境中，默认使用CPU
+        if HF_SPACE:
+            device = "cpu"  # HF Space可能没有GPU
+        
+        model = smp.create_model(
+            "DeepLabV3Plus",
+            encoder_name="efficientnet-b6",
+            in_channels=3,
+            classes=len(CLASSES),
+            encoder_weights=None
+        )
+        state_dict = torch.load(model_path, map_location=device)
+        if all(k.startswith('model.') for k in state_dict.keys()):
+            state_dict = {k[6:]: v for k, v in state_dict.items()}
+        model.load_state_dict(state_dict)
+        model.to(device)
+        model.eval()
+        print(f"Model load success: {model_path}")
+        return model
+    except Exception as e:
+        print(f"Model load fail: {e}")
+        return None
+
+# Load reference vector
+def load_reference_vector(vector_path):
+    try:
+        ref_vector = np.load(vector_path)
+        print(f"reference vector load success: {vector_path}")
+        return ref_vector
+    except Exception as e:
+        print(f"reference vector load {vector_path}: {e}")
+        return []
+
+# Load reference image
+def load_reference_images(ref_dir):
+    try:
+        image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp']
+        image_files = []
+        for ext in image_extensions:
+            image_files.extend(glob(os.path.join(ref_dir, ext)))
+        image_files.sort()
+        reference_images = []
+        for file in image_files[:4]:
+            img = cv2.imread(file)
+            if img is not None:
+                reference_images.append(img)
+        print(f"from {ref_dir} load {len(reference_images)} images")
+        return reference_images
+    except Exception as e:
+        print(f"load image failed {ref_dir}: {e}")
+        return []
+
+# Preprocess the image
+def preprocess_image(image):
+    if image.shape[2] == 4:
+        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
+    orig_h, orig_w = image.shape[:2]
+    image_resized = cv2.resize(image, (1024, 1024))
+    image_norm = image_resized.astype(np.float32) / 255.0
+    mean = np.array([0.485, 0.456, 0.406])
+    std = np.array([0.229, 0.224, 0.225])
+    image_norm = (image_norm - mean) / std
+    image_tensor = torch.from_numpy(image_norm.transpose(2, 0, 1)).float().unsqueeze(0)
+    return image_tensor, orig_h, orig_w
+
+# Generate segmentation map and visualization
+def generate_segmentation_map(prediction, orig_h, orig_w):
+    mask = prediction.argmax(1).squeeze().cpu().numpy().astype(np.uint8)
+    mask_resized = cv2.resize(mask, (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
+    kernel = np.ones((5, 5), np.uint8)
+    processed_mask = mask_resized.copy()
+    for idx in range(1, len(CLASSES)):
+        class_mask = (mask_resized == idx).astype(np.uint8)
+        dilated_mask = cv2.dilate(class_mask, kernel, iterations=2)
+        dilated_effect = dilated_mask & (mask_resized == 0)
+        processed_mask[dilated_effect > 0] = idx
+    segmentation_map = np.zeros((orig_h, orig_w, 3), dtype=np.uint8)
+    for idx, color in enumerate(COLORS):
+        segmentation_map[processed_mask == idx] = color
+    return segmentation_map
+
+# Analysis result HTML
+def create_analysis_result(mask):
+    total_pixels = mask.size
+    percentages = {cls: round((np.sum(mask == i) / total_pixels) * 100, 1)
+                   for i, cls in enumerate(CLASSES)}
+    ordered = ['sky', 'cobbles', 'plant', 'drysand', 'wetsand', 'water']
+    result = "<div style='font-size:18px;font-weight:bold;'>"
+    result += " | ".join(f"{cls}: {percentages.get(cls,0)}%" for cls in ordered)
+    result += "</div>"
+    return result
+
+# Merge and overlay
+def create_overlay(image, segmentation_map, alpha=0.5):
+    if image.shape[:2] != segmentation_map.shape[:2]:
+        segmentation_map = cv2.resize(segmentation_map, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
+    return cv2.addWeighted(image, 1-alpha, segmentation_map, alpha, 0)
+
+# Perform segmentation
+def perform_segmentation(model, image_bgr):
+    device = "cuda" if torch.cuda.is_available() and not HF_SPACE else "cpu"
+    image_rgb = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
+    image_tensor, orig_h, orig_w = preprocess_image(image_rgb)
+    with torch.no_grad():
+        prediction = model(image_tensor.to(device))
+    seg_map = generate_segmentation_map(prediction, orig_h, orig_w)  # RGB
+    overlay = create_overlay(image_rgb, seg_map)
+    mask = prediction.argmax(1).squeeze().cpu().numpy()
+    analysis = create_analysis_result(mask)
+    return seg_map, overlay, analysis
+
+# Single image processing
+def process_coastal_image(location, input_image):
+    if input_image is None:
+        return None, None, "请上传一张图片", "未检测", None
+    
+    device = "cuda" if torch.cuda.is_available() and not HF_SPACE else "cpu"
+    model = load_model(MODEL_PATHS[location], device)
+    
+    if model is None:
+        return None, None, f"错误：无法加载模型", "未检测", None
+    
+    ref_vector = load_reference_vector(REFERENCE_VECTOR_PATHS[location]) if os.path.exists(REFERENCE_VECTOR_PATHS[location]) else []
+    ref_images = load_reference_images(REFERENCE_IMAGE_DIRS[location])
+    
+    outlier_status = "未检测"
+    is_outlier = False
+    image_bgr = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
+    
+    if len(ref_vector) > 0:
+        filtered, _ = detect_outliers(ref_images, [image_bgr], ref_vector)
+        is_outlier = len(filtered) == 0
+    else:
+        filtered, _ = detect_outliers(ref_images, [image_bgr])
+        is_outlier = len(filtered) == 0
+    
+    outlier_status = "异常检测: <span style='color:red;font-weight:bold'>未通过</span>" if is_outlier else "异常检测: <span style='color:green;font-weight:bold'>通过</span>"
+    seg_map, overlay, analysis = perform_segmentation(model, image_bgr)
+    
+    # 在HF环境中不上传，只返回本地结果
+    url = "本地存储"
+    if not HF_SPACE:
+        try:
+            url = upload_mask(Image.fromarray(seg_map), prefix=location.replace(' ', '_'))
+        except Exception as e:
+            print(f"Upload failed: {e}")
+            url = "上传错误"
+    
+    if is_outlier:
+        analysis = "<div style='color:red;font-weight:bold;margin-bottom:10px'>警告：图像未通过异常检测，结果可能不准确！</div>" + analysis
+    
+    return seg_map, overlay, analysis, outlier_status, url
+
+# Spacial Alignment
+def process_with_alignment(location, reference_image, input_image):
+    if reference_image is None or input_image is None:
+        return None, None, None, None, "请上传参考图像和需要分析的图像", "未处理", None
+    
+    device = "cuda" if torch.cuda.is_available() and not HF_SPACE else "cpu"
+    model = load_model(MODEL_PATHS[location], device)
+    
+    if model is None:
+        return None, None, None, None, "错误：无法加载模型", "未处理", None
+    
+    ref_bgr = cv2.cvtColor(np.array(reference_image), cv2.COLOR_RGB2BGR)
+    tgt_bgr = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
+    
+    aligned, _ = align_images([ref_bgr, tgt_bgr], [np.zeros_like(ref_bgr), np.zeros_like(tgt_bgr)])
+    aligned_tgt_bgr = aligned[1]
+    
+    seg_map, overlay, analysis = perform_segmentation(model, aligned_tgt_bgr)
+    
+    # 在HF环境中不上传，只返回本地结果
+    url = "本地存储"
+    if not HF_SPACE:
+        try:
+            url = upload_mask(Image.fromarray(seg_map), prefix="aligned_" + location.replace(' ', '_'))
+        except Exception as e:
+            print(f"Upload failed: {e}")
+            url = "上传错误"
+    
+    status = "空间对齐: <span style='color:green;font-weight:bold'>完成</span>"
+    ref_rgb = cv2.cvtColor(ref_bgr, cv2.COLOR_BGR2RGB)
+    aligned_tgt_rgb = cv2.cvtColor(aligned_tgt_bgr, cv2.COLOR_BGR2RGB)
+    
+    return ref_rgb, aligned_tgt_rgb, seg_map, overlay, analysis, status, url
+
+# Create the Gradio interface
+def create_interface():
+    scale = 0.5
+    disp_w, disp_h = int(1365*scale), int(1024*scale)
+    with gr.Blocks(title="海岸侵蚀分析系统") as demo:
+        gr.Markdown("""# 海岸侵蚀分析系统
+
+上传海岸照片进行分析，包括分割和空间对齐功能。""")
+        with gr.Tabs():
+            with gr.TabItem("单张图像分割"):
+                with gr.Row():
+                    loc1 = gr.Radio(list(MODEL_PATHS.keys()), label="选择模型", value=list(MODEL_PATHS.keys())[0])
+                with gr.Row():
+                    inp = gr.Image(label="输入图像", type="numpy", image_mode="RGB")
+                    seg = gr.Image(label="分割图像", type="numpy", width=disp_w, height=disp_h)
+                    ovl = gr.Image(label="叠加图像", type="numpy", width=disp_w, height=disp_h)
+                with gr.Row():
+                    btn1 = gr.Button("执行分割")
+                url1 = gr.Text(label="分割图URL")
+                status1 = gr.HTML(label="异常检测状态")
+                res1 = gr.HTML(label="分析结果")
+                btn1.click(fn=process_coastal_image,inputs=[loc1, inp],outputs=[seg, ovl, res1, status1, url1])
+            
+            with gr.TabItem("空间对齐分割"):
+                with gr.Row():
+                    loc2 = gr.Radio(list(MODEL_PATHS.keys()), label="选择模型", value=list(MODEL_PATHS.keys())[0])
+                with gr.Row():
+                    ref_img = gr.Image(label="参考图像", type="numpy", image_mode="RGB")
+                    tgt_img = gr.Image(label="待分析图像", type="numpy", image_mode="RGB")
+                with gr.Row():
+                    btn2 = gr.Button("执行空间对齐分割")
+                with gr.Row():
+                    orig = gr.Image(label="原始图像", type="numpy", width=disp_w, height=disp_h)
+                    aligned = gr.Image(label="对齐后图像", type="numpy", width=disp_w, height=disp_h)
+                with gr.Row():
+                    seg2 = gr.Image(label="分割图像", type="numpy", width=disp_w, height=disp_h)
+                    ovl2 = gr.Image(label="叠加图像", type="numpy", width=disp_w, height=disp_h)
+                    url2 = gr.Text(label="分割图URL")
+                status2 = gr.HTML(label="空间对齐状态")
+                res2 = gr.HTML(label="分析结果")
+                btn2.click(fn=process_with_alignment, inputs=[loc2, ref_img, tgt_img], outputs=[orig, aligned, seg2, ovl2, res2, status2, url2])
+    return demo
+
+if __name__ == "__main__":
+    for path in ["models", "reference_images/MM", "reference_images/SJ"]:
+        os.makedirs(path, exist_ok=True)
+    for p in MODEL_PATHS.values():
+        if not os.path.exists(p):
+            print(f"警告：模型文件 {p} 不存在！")
+    demo = create_interface()
+    # 在HF环境中使用适当的启动配置
+    if HF_SPACE:
+        demo.launch()
+    else:
+        demo.launch(share=True)