init

app.py

@@ -5,6 +5,9 @@ import torch
 import gradio as gr
 import segmentation_models_pytorch as smp
 from PIL import Image
+import boto3
+import uuid
+import io
 from glob import glob
 from pipeline.ImgOutlier import detect_outliers
 from pipeline.normalization import align_images
@@ -12,13 +15,60 @@ from pipeline.normalization import align_images
 # 检测是否在Hugging Face环境中运行
 HF_SPACE = os.environ.get('SPACE_ID') is not None
 
-# 尝试导入上传模块，如果不在HF环境中才需要
-if not HF_SPACE:
+# DigitalOcean Spaces上传函数
+def upload_mask(image, prefix="mask"):
+    """
+    将分割掩码图像上传到DigitalOcean Spaces
+    
+    Args:
+        image: PIL Image对象
+        prefix: 文件名前缀
+        
+    Returns:
+        上传文件的URL
+    """
     try:
-        from uploader.do_spaces import upload_mask
-    except ImportError:
-        def upload_mask(image, prefix=""):
-            return "上传模块未加载"
+        # 从环境变量获取凭据
+        do_key = os.environ.get('DO_SPACES_KEY')
+        do_secret = os.environ.get('DO_SPACES_SECRET')
+        do_region = os.environ.get('DO_SPACES_REGION')
+        do_bucket = os.environ.get('DO_SPACES_BUCKET')
+        
+        # 校验凭据是否存在
+        if not all([do_key, do_secret, do_region, do_bucket]):
+            return "DigitalOcean凭据未设置"
+        
+        # 创建S3客户端
+        session = boto3.session.Session()
+        client = session.client('s3',
+                               region_name=do_region,
+                               endpoint_url=f'https://{do_region}.digitaloceanspaces.com',
+                               aws_access_key_id=do_key,
+                               aws_secret_access_key=do_secret)
+        
+        # 生成唯一文件名
+        filename = f"{prefix}_{uuid.uuid4().hex}.png"
+        
+        # 将图像转换为字节流
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0)
+        
+        # 上传到Spaces
+        client.upload_fileobj(
+            img_byte_arr,
+            do_bucket,
+            filename,
+            ExtraArgs={'ACL': 'public-read', 'ContentType': 'image/png'}
+        )
+        
+        # 返回公共URL
+        url = f'https://{do_bucket}.{do_region}.digitaloceanspaces.com/{filename}'
+        return url
+    
+    except Exception as e:
+        print(f"上传失败: {str(e)}")
+        return f"上传错误: {str(e)}"
 
 # Global Configuration
 MODEL_PATHS = {
@@ -54,6 +104,8 @@ def load_model(model_path, device="cuda"):
         # 如果在HF环境中，默认使用CPU
         if HF_SPACE:
             device = "cpu"  # HF Space可能没有GPU
+        elif not torch.cuda.is_available():
+            device = "cpu"  # 本地环境也可能没有GPU
         
         model = smp.create_model(
             "DeepLabV3Plus",
@@ -68,25 +120,33 @@ def load_model(model_path, device="cuda"):
         model.load_state_dict(state_dict)
         model.to(device)
         model.eval()
-        print(f"Model load success: {model_path}")
+        print(f"模型加载成功: {model_path}")
         return model
     except Exception as e:
-        print(f"Model load fail: {e}")
+        print(f"模型加载失败: {e}")
         return None
 
 # Load reference vector
 def load_reference_vector(vector_path):
     try:
+        if not os.path.exists(vector_path):
+            print(f"参考向量文件不存在: {vector_path}")
+            return []
         ref_vector = np.load(vector_path)
-        print(f"reference vector load success: {vector_path}")
+        print(f"参考向量加载成功: {vector_path}")
         return ref_vector
     except Exception as e:
-        print(f"reference vector load {vector_path}: {e}")
+        print(f"参考向量加载失败 {vector_path}: {e}")
         return []
 
 # Load reference image
 def load_reference_images(ref_dir):
     try:
+        if not os.path.exists(ref_dir):
+            print(f"参考图像目录不存在: {ref_dir}")
+            os.makedirs(ref_dir, exist_ok=True)
+            return []
+            
         image_extensions = ['*.jpg', '*.jpeg', '*.png', '*.bmp']
         image_files = []
         for ext in image_extensions:
@@ -97,10 +157,10 @@ def load_reference_images(ref_dir):
             img = cv2.imread(file)
             if img is not None:
                 reference_images.append(img)
-        print(f"from {ref_dir} load {len(reference_images)} images")
+        print(f"从 {ref_dir} 加载了 {len(reference_images)} 张图像")
         return reference_images
     except Exception as e:
-        print(f"load image failed {ref_dir}: {e}")
+        print(f"加载图像失败 {ref_dir}: {e}")
         return []
 
 # Preprocess the image
@@ -173,7 +233,7 @@ def process_coastal_image(location, input_image):
     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_vector = load_reference_vector(REFERENCE_VECTOR_PATHS[location])
     ref_images = load_reference_images(REFERENCE_IMAGE_DIRS[location])
     
     outlier_status = "未检测"
@@ -183,21 +243,23 @@ def process_coastal_image(location, input_image):
     if len(ref_vector) > 0:
         filtered, _ = detect_outliers(ref_images, [image_bgr], ref_vector)
         is_outlier = len(filtered) == 0
-    else:
+    elif len(ref_images) > 0:
         filtered, _ = detect_outliers(ref_images, [image_bgr])
         is_outlier = len(filtered) == 0
+    else:
+        print("警告：没有参考图像或参考向量可用于异常检测")
+        is_outlier = False
     
     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环境中不上传，只返回本地结果
+    # 尝试上传到DigitalOcean Spaces
     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 = "上传错误"
+    try:
+        url = upload_mask(Image.fromarray(seg_map), prefix=location.replace(' ', '_'))
+    except Exception as e:
+        print(f"上传失败: {e}")
+        url = f"上传错误: {str(e)}"
     
     if is_outlier:
         analysis = "<div style='color:red;font-weight:bold;margin-bottom:10px'>警告：图像未通过异常检测，结果可能不准确！</div>" + analysis
@@ -218,19 +280,22 @@ def process_with_alignment(location, reference_image, input_image):
     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]
+    try:
+        aligned, _ = align_images([ref_bgr, tgt_bgr], [np.zeros_like(ref_bgr), np.zeros_like(tgt_bgr)])
+        aligned_tgt_bgr = aligned[1]
+    except Exception as e:
+        print(f"空间对齐失败: {e}")
+        return None, None, None, None, f"空间对齐失败: {str(e)}", "处理失败", None
     
     seg_map, overlay, analysis = perform_segmentation(model, aligned_tgt_bgr)
     
-    # 在HF环境中不上传，只返回本地结果
+    # 尝试上传到DigitalOcean Spaces
     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 = "上传错误"
+    try:
+        url = upload_mask(Image.fromarray(seg_map), prefix="aligned_" + location.replace(' ', '_'))
+    except Exception as e:
+        print(f"上传失败: {e}")
+        url = f"上传错误: {str(e)}"
     
     status = "空间对齐: <span style='color:green;font-weight:bold'>完成</span>"
     ref_rgb = cv2.cvtColor(ref_bgr, cv2.COLOR_BGR2RGB)
@@ -259,7 +324,7 @@ def create_interface():
                 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])
+                btn1.click(fn=process_coastal_image, inputs=[loc1, inp], outputs=[seg, ovl, res1, status1, url1])
             
             with gr.TabItem("空间对齐分割"):
                 with gr.Row():
@@ -275,18 +340,33 @@ def create_interface():
                 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")
+                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} 不存在！")
+    
+    # 检查DigitalOcean凭据是否存在
+    do_creds = [
+        os.environ.get('DO_SPACES_KEY'),
+        os.environ.get('DO_SPACES_SECRET'),
+        os.environ.get('DO_SPACES_REGION'),
+        os.environ.get('DO_SPACES_BUCKET')
+    ]
+    if not all(do_creds):
+        print("警告：DigitalOcean Spaces凭据不完整，上传功能可能不可用")
+    
+    # 创建并启动界面
     demo = create_interface()
     # 在HF环境中使用适当的启动配置
     if HF_SPACE:

pipeline/HSV.py

@@ -0,0 +1,16 @@
+import cv2
+import numpy as np
+
+def preprocess_images(images, V_FIXED = 200):
+    fixed_images = []
+    for image in images:
+        hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
+
+        hsv_fixed = hsv.copy()
+        hsv_fixed[:, :, 2] = (hsv[:, :, 2] / hsv[:, :, 2].max()) * V_FIXED
+        hsv_fixed[:, :, 1] = hsv_fixed[:, :, 1] * (hsv_fixed[:, :, 2] / hsv[:, :, 2].max())
+        hsv_fixed[:, :, 1] = np.clip(hsv_fixed[:, :, 1], 0, 255)
+
+        fixed_image = cv2.cvtColor(hsv_fixed, cv2.COLOR_HSV2BGR)
+        fixed_images.append(fixed_image)
+    return fixed_images

pipeline/ImgOutlier.py

@@ -0,0 +1,200 @@
+import numpy as np
+import torch
+import torchvision
+from PIL import Image
+from torch import nn
+from torchvision import transforms as tr
+from torchvision.models import vit_h_14
+import cv2
+
+class CosineSimilarity:
+    def __init__(self, vector='feature', threshold=0.8, mean_vec=[], device=None):
+        """
+        Initialize the CosineSimilarity class.
+
+        Args:
+            vector (str): Type of vector to use ('feature' or 'image')
+            threshold (float): Threshold for determining outliers
+            mean_vec (numpy vector): Preloaded reference vector for comparison
+            device (str): Device to use for computation (default: 'mps' if available, else 'cuda' if available, else 'cpu')
+        """
+        if device is None:
+            if torch.backends.mps.is_available():
+                self.device = 'mps'
+            elif torch.cuda.is_available():
+                self.device = 'cuda'
+            else:
+                self.device = 'cpu'
+        else:
+            self.device = device
+
+        self.vector = vector
+        self.threshold = threshold
+        self.model_instance = None
+        self.mean_vec = mean_vec
+
+    def model(self):
+        """Initialize and return the ViT model."""
+        if self.model_instance is None:
+            wt = torchvision.models.ViT_H_14_Weights.DEFAULT
+            self.model_instance = vit_h_14(weights=wt)
+            self.model_instance.heads = nn.Sequential(*list(self.model_instance.heads.children())[:-1])
+            self.model_instance = self.model_instance.to(self.device)
+        return self.model_instance
+
+    def process_image(self, cv2_img):
+        """
+        Process a cv2 image for the model.
+
+        Args:
+            cv2_img: OpenCV image (BGR format)
+
+        Returns:
+            Processed tensor
+        """
+        # Convert BGR to RGB
+        rgb_img = cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB)
+        # Convert to PIL Image
+        pil_img = Image.fromarray(rgb_img)
+
+        # A set of transformations to prepare the image in tensor format
+        transformations = tr.Compose([
+            tr.ToTensor(),
+            tr.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
+            tr.Resize((518, 518))
+        ])
+
+        # preparing the image
+        img_tensor = transformations(pil_img).float()
+
+        if self.vector == 'image':
+            img_tensor = img_tensor.flatten()
+
+        img_tensor = img_tensor.unsqueeze_(0)
+
+        if self.vector == 'feature':
+            img_tensor = img_tensor.to(self.device)
+
+        return img_tensor
+
+    def get_embeddings(self, ref_images, test_images):
+        """
+        Get embeddings for reference and test images.
+
+        Args:
+            ref_images: List of cv2 reference images
+            test_images: List of cv2 test images
+
+        Returns:
+            Reference embedding, list of test embeddings
+        """
+        model = self.model()
+
+        # Process test images
+        emb_test = []
+        for img in test_images:
+            processed_img = self.process_image(img)
+            if self.vector == 'feature':
+                emb = model(processed_img).detach().cpu()
+                emb_test.append(emb)
+            else:  # 'image'
+                emb_test.append(processed_img)
+
+        # This checks if a reference vector is loaded, if so the process of getting
+        # reference embeddings can be skipped for efficiency
+        if len(self.mean_vec) > 0:
+            emb_ref = torch.tensor(self.mean_vec)
+
+        # Process reference images if necessary
+        else:
+            if self.vector == 'feature':
+                # Standard method of getting reference embedding vector
+                emb_ref_list = []
+                for img in ref_images:
+                    processed_img = self.process_image(img)
+                    emb = model(processed_img).detach().cpu()
+                    emb_ref_list.append(emb)
+
+                # Average the reference embeddings
+                emb_ref = torch.mean(torch.stack(emb_ref_list), dim=0)
+            
+            else:  # 'image'
+                emb_ref_list = []
+                for img in ref_images:
+                    processed_img = self.process_image(img)
+                    emb_ref_list.append(processed_img)
+
+                # Average the reference images
+                emb_ref = torch.mean(torch.stack(emb_ref_list), dim=0)
+
+        return emb_ref, emb_test
+
+    def find_outliers(self, ref_images, test_images):
+        """
+        Find outliers in test images compared to reference images.
+
+        Args:
+            ref_images: List of cv2 reference images
+            test_images: List of cv2 test images
+
+        Returns:
+            mask: Boolean array where True indicates an outlier
+            scores: Similarity scores for each test image
+        """
+        emb_ref, emb_test = self.get_embeddings(ref_images, test_images)
+
+        scores = []
+        mask = []
+
+        for i in range(len(emb_test)):
+            score = torch.nn.functional.cosine_similarity(emb_ref, emb_test[i])
+            score_value = score.item()
+            scores.append(round(score_value, 4))
+            # True if it's an outlier (below threshold)
+            mask.append(score_value <= self.threshold)
+
+        return np.array(mask), scores, emb_ref
+
+    def filter_outliers(self, ref_images, test_images):
+        """
+        Filter out outliers from test images.
+
+        Args:
+            ref_images: List of cv2 reference images
+            test_images: List of cv2 test images
+
+        Returns:
+            filtered_images: List of non-outlier test images
+            outlier_mask: Boolean array where True indicates an outlier
+            scores: Similarity scores for each test image
+        """
+        outlier_mask, scores, mean = self.find_outliers(ref_images, test_images)
+
+        # Filter out outliers (keep only non-outliers)
+        filtered_images = [img for i, img in enumerate(test_images) if not outlier_mask[i]]
+
+        return filtered_images, outlier_mask, scores, mean
+
+def detect_outliers(ref_imgs, imgs, mean_vec=[]):
+    """
+    Detects outliers in a set of test images, can use a reference vector
+
+    Args:
+        ref_images: List of cv2 reference images
+        images: List of cv2 test images
+        mean_vec: optional pre-computed reference vector
+
+    Returns:
+        filtered_images: List of non-outlier test images
+        mean: the reference vector used (if a new reference vector should be saved)
+    """
+    
+    similarity = CosineSimilarity(vector='feature', threshold=0.8, mean_vec=mean_vec)
+
+    # Get outlier mask, scores, and reference vector
+    outlier_mask, scores, mean_vector = similarity.find_outliers(ref_imgs, imgs)
+
+    # Filter out outliers
+    filtered_images = [img for i, img in enumerate(imgs) if not outlier_mask[i]]
+
+    return filtered_images, mean_vector

pipeline/normalization.py

@@ -0,0 +1,77 @@
+from transformers import AutoImageProcessor, AutoModel
+import torch
+import numpy as np
+import cv2
+
+
+def align_images(images, segs):
+    """
+    Align images using SuperGlue for feature matching.
+
+    Args:
+        images: List of input images
+        segs: List of segmentation images
+
+    Returns:
+        Tuple of (aligned images, aligned segmentation images)
+    """
+    if not images or len(images) < 2:
+        return images, segs
+
+    reference = images[0]
+    reference_seg = segs[0]
+    aligned_images = [reference]
+    aligned_images_seg = [reference_seg]
+
+    # Load SuperGlue model and processor
+    processor = AutoImageProcessor.from_pretrained("magic-leap-community/superglue_outdoor")
+    model = AutoModel.from_pretrained("magic-leap-community/superglue_outdoor")
+
+    for i in range(1, len(images)):
+        current = images[i]
+        current_seg = segs[i]
+
+        # Process image pair
+        image_pair = [reference, current]
+        inputs = processor(image_pair, return_tensors="pt")
+
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        # Get matches
+        image_sizes = [[(img.shape[0], img.shape[1]) for img in image_pair]]
+        matches = processor.post_process_keypoint_matching(outputs, image_sizes, threshold=0.2)
+
+        # Extract matching keypoints
+        match_data = matches[0]
+        keypoints0 = match_data["keypoints0"].numpy()
+        keypoints1 = match_data["keypoints1"].numpy()
+
+        # Filter matches by confidence
+        valid_matches = match_data["matching_scores"] > 0.5
+        if sum(valid_matches) < 4:
+            print(f"Not enough confident matches for image {i}, keeping original")
+            aligned_images.append(current)
+            aligned_images_seg.append(current_seg)
+            continue
+
+        # Get matching points
+        src_pts = keypoints1[valid_matches].reshape(-1, 1, 2)
+        dst_pts = keypoints0[valid_matches].reshape(-1, 1, 2)
+
+        # Find homography
+        H, mask = cv2.findHomography(src_pts, dst_pts, cv2.RANSAC, 5.0)
+
+        if H is not None:
+            # Apply homography
+            h, w = reference.shape[:2]
+            aligned = cv2.warpPerspective(current, H, (w, h))
+            aligned_images.append(aligned)
+            aligned_seg = cv2.warpPerspective(current_seg, H, (w, h))
+            aligned_images_seg.append(aligned_seg)
+        else:
+            print(f"Could not find homography for image {i}, keeping original")
+            aligned_images.append(current)
+            aligned_images_seg.append(current_seg)
+
+    return aligned_images, aligned_images_seg

requirements.txt

@@ -0,0 +1,9 @@
+torch
+torchvision
+segmentation_models_pytorch
+gradio
+opencv-python
+numpy
+pillow
+transformers
+boto3

uploader/do_spaces.py

@@ -0,0 +1,60 @@
+import os
+import boto3
+from botocore.client import Config
+import uuid
+from PIL import Image
+import io
+
+def upload_mask(image, prefix="mask"):
+    """
+    将分割掩码图像上传到DigitalOcean Spaces
+    
+    Args:
+        image: PIL Image对象
+        prefix: 文件名前缀
+        
+    Returns:
+        上传文件的URL
+    """
+    try:
+        # 从环境变量获取凭据
+        do_key = os.environ.get('DO_SPACES_KEY')
+        do_secret = os.environ.get('DO_SPACES_SECRET')
+        do_region = os.environ.get('DO_SPACES_REGION')
+        do_bucket = os.environ.get('DO_SPACES_BUCKET')
+        
+        # 校验凭据是否存在
+        if not all([do_key, do_secret, do_region, do_bucket]):
+            raise ValueError("缺少DigitalOcean Spaces凭据")
+        
+        # 创建S3客户端
+        session = boto3.session.Session()
+        client = session.client('s3',
+                               region_name=do_region,
+                               endpoint_url=f'https://{do_region}.digitaloceanspaces.com',
+                               aws_access_key_id=do_key,
+                               aws_secret_access_key=do_secret)
+        
+        # 生成唯一文件名
+        filename = f"{prefix}_{uuid.uuid4().hex}.png"
+        
+        # 将图像转换为字节流
+        img_byte_arr = io.BytesIO()
+        image.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0)
+        
+        # 上传到Spaces
+        client.upload_fileobj(
+            img_byte_arr,
+            do_bucket,
+            filename,
+            ExtraArgs={'ACL': 'public-read', 'ContentType': 'image/png'}
+        )
+        
+        # 返回公共URL
+        url = f'https://{do_bucket}.{do_region}.digitaloceanspaces.com/{filename}'
+        return url
+    
+    except Exception as e:
+        print(f"上传失败: {str(e)}")
+        return f"上传错误: {str(e)}"