Upload app.py

app.py

@@ -89,12 +89,14 @@ high_level_config = {
     "resolution": 518,
 }
 
-# GroundingDINO and SAM Configuration
-GROUNDING_DINO_MODEL_ID = "IDEA-Research/grounding-dino-tiny"
+# GroundingDINO and SAM Configuration (CPU-friendly versions)
+GROUNDING_DINO_MODEL_ID = "IDEA-Research/grounding-dino-tiny"  # 已经是tiny版本
 GROUNDING_DINO_BOX_THRESHOLD = 0.25
 GROUNDING_DINO_TEXT_THRESHOLD = 0.2
 
-SAM_MODEL_ID = "facebook/sam-vit-huge"
+# 使用 MobileSAM (CPU友好，比SAM-huge快60倍，只有10MB)
+SAM_MODEL_ID = "dhkim2810/MobileSAM"  # 轻量级SAM，适合CPU
+USE_MOBILE_SAM = True  # 标记使用MobileSAM
 
 DEFAULT_TEXT_PROMPT = "chair . table . sofa . bed . desk . cabinet"
 
@@ -154,8 +156,8 @@ sam_predictor = None
 # Model Loading Functions
 # ============================================================================
 
-def load_grounding_dino_model(device):
-    """Load GroundingDINO model from HuggingFace"""
+def load_grounding_dino_model(device="cpu"):
+    """Load GroundingDINO model from HuggingFace (CPU优化)"""
     global grounding_dino_model, grounding_dino_processor
     
     if grounding_dino_model is not None:
@@ -165,13 +167,15 @@ def load_grounding_dino_model(device):
     try:
         from transformers import AutoProcessor, AutoModelForZeroShotObjectDetection
         
-        print(f"📥 Loading GroundingDINO from HuggingFace: {GROUNDING_DINO_MODEL_ID}")
+        # 强制使用 CPU 进行分割（节省 GPU 资源）
+        seg_device = "cpu"
+        print(f"📥 Loading GroundingDINO from HuggingFace: {GROUNDING_DINO_MODEL_ID} (使用 {seg_device.upper()})")
         grounding_dino_processor = AutoProcessor.from_pretrained(GROUNDING_DINO_MODEL_ID)
         grounding_dino_model = AutoModelForZeroShotObjectDetection.from_pretrained(
             GROUNDING_DINO_MODEL_ID
-        ).to(device).eval()
+        ).to(seg_device).eval()
         
-        print("✅ GroundingDINO loaded successfully")
+        print(f"✅ GroundingDINO loaded successfully on {seg_device.upper()}")
         
     except Exception as e:
         print(f"❌ GroundingDINO loading failed: {e}")
@@ -179,8 +183,8 @@ def load_grounding_dino_model(device):
         traceback.print_exc()
 
 
-def load_sam_model(device):
-    """Load SAM model from HuggingFace"""
+def load_sam_model(device="cpu"):
+    """Load MobileSAM model from HuggingFace (CPU优化，比SAM快60倍)"""
     global sam_predictor
     
     if sam_predictor is not None:
@@ -190,8 +194,12 @@ def load_sam_model(device):
     try:
         from transformers import SamModel, SamProcessor
         
-        print(f"📥 Loading SAM from HuggingFace: {SAM_MODEL_ID}")
-        sam_model = SamModel.from_pretrained(SAM_MODEL_ID).to(device).eval()
+        # 强制使用 CPU 进行分割（MobileSAM 专为移动设备/CPU优化）
+        seg_device = "cpu"
+        print(f"📥 Loading MobileSAM from HuggingFace: {SAM_MODEL_ID} (使用 {seg_device.upper()})")
+        print(f"   💡 MobileSAM 是轻量级版本，比 SAM-huge 快60倍，只有10MB，适合CPU运行")
+        
+        sam_model = SamModel.from_pretrained(SAM_MODEL_ID).to(seg_device).eval()
         sam_processor = SamProcessor.from_pretrained(SAM_MODEL_ID)
         
         # Wrap in a predictor-like interface
@@ -210,20 +218,22 @@ def load_sam_model(device):
                     self.image = Image.fromarray((image * 255).astype(np.uint8))
             
             def predict(self, box, multimask_output=False):
-                """Predict mask from box"""
+                """Predict mask from box (CPU优化)"""
                 inputs = self.processor(
                     self.image, 
                     input_boxes=[[[box]]], 
                     return_tensors="pt"
-                ).to(self.device)
+                )
+                # 确保在CPU上运行
+                inputs = {k: v.to(self.device) for k, v in inputs.items() if isinstance(v, torch.Tensor)}
                 
                 with torch.no_grad():
                     outputs = self.model(**inputs)
                 
                 masks = self.processor.image_processor.post_process_masks(
                     outputs.pred_masks.cpu(),
-                    inputs["original_sizes"].cpu(),
-                    inputs["reshaped_input_sizes"].cpu()
+                    inputs["original_sizes"].cpu() if "original_sizes" in inputs else outputs.pred_masks.new_tensor([[self.image.height, self.image.width]]),
+                    inputs["reshaped_input_sizes"].cpu() if "reshaped_input_sizes" in inputs else outputs.pred_masks.new_tensor([[self.image.height, self.image.width]])
                 )[0].squeeze().numpy()
                 
                 if len(masks.shape) == 2:
@@ -231,8 +241,8 @@ def load_sam_model(device):
                 
                 return masks, None, None
         
-        sam_predictor = SAMPredictor(sam_model, sam_processor, device)
-        print("✅ SAM loaded successfully")
+        sam_predictor = SAMPredictor(sam_model, sam_processor, seg_device)
+        print(f"✅ MobileSAM loaded successfully on {seg_device.upper()}")
         
     except Exception as e:
         print(f"❌ SAM loading failed: {e}")
@@ -261,14 +271,14 @@ def generate_distinct_colors(n):
     return colors
 
 
-def run_grounding_dino_detection(image_np, text_prompt, device):
-    """Run GroundingDINO detection"""
+def run_grounding_dino_detection(image_np, text_prompt, device="cpu"):
+    """Run GroundingDINO detection (CPU优化)"""
     if grounding_dino_model is None or grounding_dino_processor is None:
         print("⚠️ GroundingDINO not loaded")
         return []
     
     try:
-        print(f"🔍 GroundingDINO detection: {text_prompt}")
+        print(f"🔍 GroundingDINO detection (CPU): {text_prompt}")
         
         # Convert to PIL Image
         if image_np.dtype == np.uint8:
@@ -276,9 +286,10 @@ def run_grounding_dino_detection(image_np, text_prompt, device):
         else:
             pil_image = Image.fromarray((image_np * 255).astype(np.uint8))
         
-        # Preprocess
+        # Preprocess - 强制使用CPU
+        seg_device = "cpu"
         inputs = grounding_dino_processor(images=pil_image, text=text_prompt, return_tensors="pt")
-        inputs = {k: v.to(device) for k, v in inputs.items()}
+        inputs = {k: v.to(seg_device) for k, v in inputs.items() if isinstance(v, torch.Tensor)}
         
         # Inference
         with torch.no_grad():
@@ -705,11 +716,11 @@ def run_model(
 
     model.eval()
 
-    # Load segmentation models if enabled
+    # Load segmentation models if enabled (使用CPU节省GPU资源)
     if enable_segmentation:
-        progress(0.1, desc="🎯 加载分割模型...")
-        load_grounding_dino_model(device)
-        load_sam_model(device)
+        progress(0.1, desc="🎯 加载分割模型 (CPU)...")
+        load_grounding_dino_model("cpu")  # 分割使用CPU
+        load_sam_model("cpu")  # MobileSAM在CPU上运行良好
 
     # Load images
     progress(0.15, desc="📷 加载图片...")
@@ -800,7 +811,7 @@ def run_model(
             else:
                 ref_image_np = ref_image
 
-            detections = run_grounding_dino_detection(ref_image_np, text_prompt, device)
+            detections = run_grounding_dino_detection(ref_image_np, text_prompt, "cpu")  # 使用CPU进行检测
 
             if len(detections) > 0:
                 boxes = [d['bbox'] for d in detections]
@@ -820,6 +831,9 @@ def run_model(
                 all_view_masks.append([])
 
         # Match objects across views
+        total_detections = sum(len(dets) for dets in all_view_detections)
+        print(f"\n📊 总检测数: {total_detections}")
+        
         if any(len(dets) > 0 for dets in all_view_detections):
             progress(0.85, desc="🔗 匹配跨视图物体...")
             object_id_map, unique_objects = match_objects_across_views(all_view_detections)
@@ -830,6 +844,13 @@ def run_model(
                 processed_data, all_view_detections, all_view_masks,
                 object_id_map, unique_objects, target_dir
             )
+            
+            if segmented_glb:
+                print(f"✅ 分割3D模型已生成: {segmented_glb}")
+            else:
+                print(f"⚠️ 分割3D模型生成失败")
+        else:
+            print(f"⚠️ 未检测到任何物体，无法生成分割模型")
 
     # Cleanup
     progress(0.95, desc="🧹 清理内存...")
@@ -1174,6 +1195,13 @@ def gradio_demo(
     )
     
     progress(1.0, desc="✅ 全部完成！")
+    
+    # 添加分割状态信息
+    if enable_segmentation:
+        if segmented_glb:
+            log_msg += f"\n🎨 分割模型已生成"
+        else:
+            log_msg += f"\n⚠️ 未检测到物体，无分割模型"
 
     return (
         glbfile,
@@ -1652,6 +1680,37 @@ CUSTOM_CSS = GRADIO_CSS + """
 .tab-content {
     min-height: 550px !important;
 }
+
+/* 增强文件上传区域 */
+.file-upload-enhanced {
+    position: relative;
+}
+"""
+
+# JavaScript for paste support
+PASTE_JS = """
+<script>
+// 添加粘贴板支持
+document.addEventListener('paste', function(e) {
+    const items = e.clipboardData.items;
+    for (let i = 0; i < items.length; i++) {
+        if (items[i].type.indexOf('image') !== -1) {
+            const blob = items[i].getAsFile();
+            const fileInput = document.querySelector('input[type="file"][multiple]');
+            if (fileInput) {
+                const dataTransfer = new DataTransfer();
+                dataTransfer.items.add(blob);
+                fileInput.files = dataTransfer.files;
+                fileInput.dispatchEvent(new Event('change', { bubbles: true }));
+                console.log('✅ 图片已从剪贴板粘贴');
+            }
+        }
+    }
+});
+
+// 添加提示信息
+console.log('💡 粘贴板功能已启用：使用 Ctrl+V 可直接粘贴截图');
+</script>
 """
 
 with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="MapAnything V2 - 3D重建与物体分割") as demo:
@@ -1659,6 +1718,9 @@ with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="MapAnything V2 - 3D重建与
     processed_data_state = gr.State(value=None)
     measure_points_state = gr.State(value=[])
 
+    # 添加粘贴板支持的 JavaScript
+    gr.HTML(PASTE_JS)
+
     # 顶部标题
     gr.HTML("""
     <div style="text-align: center; margin: 20px 0;">
@@ -1674,13 +1736,23 @@ with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="MapAnything V2 - 3D重建与
         with gr.Column(scale=1, min_width=300):
             gr.Markdown("### 📤 输入")
             
+            # 统一上传组件（支持文件、拖拽、粘贴板）
             unified_upload = gr.File(
                 file_count="multiple",
-                label="上传视频或图片",
+                label="上传视频或图片（支持拖拽、粘贴Ctrl+V📋）",
                 interactive=True,
                 file_types=["image", "video"],
             )
             
+            # 摄像头输入（折叠式）
+            with gr.Accordion("📷 使用摄像头拍照", open=False):
+                camera_input = gr.Image(
+                    label="拍照后自动添加",
+                    sources=["webcam"],
+                    type="filepath",
+                    interactive=True,
+                )
+            
             with gr.Row():
                 s_time_interval = gr.Slider(
                     minimum=0.1, maximum=5.0, value=1.0, step=0.1,
@@ -1701,9 +1773,8 @@ with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="MapAnything V2 - 3D重建与
                 show_download_button=True, object_fit="contain", preview=True
             )
             
-            
             clear_uploads_btn = gr.ClearButton(
-                [unified_upload, image_gallery],
+                [unified_upload, camera_input, image_gallery],
                 value="清空上传",
                 variant="secondary",
                 size="sm",
@@ -1712,7 +1783,7 @@ with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="MapAnything V2 - 3D重建与
             with gr.Row():
                 submit_btn = gr.Button("🚀 开始重建", variant="primary", scale=2)
                 clear_btn = gr.ClearButton(
-                    [unified_upload, target_dir_output, image_gallery],
+                    [unified_upload, camera_input, target_dir_output, image_gallery],
                     value="🗑️ 清空", scale=1
                 )
 
@@ -1803,8 +1874,9 @@ with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="MapAnything V2 - 3D重建与
                 )
                 
                 gr.Markdown("#### 分割参数")
+                gr.Markdown("💡 **说明**: 分割使用 CPU 运行（MobileSAM轻量级模型），不占用GPU资源")
                 enable_segmentation = gr.Checkbox(
-                    label="启用语义分割", value=False
+                    label="启用语义分割 (CPU)", value=False
                 )
                 
                 text_prompt = gr.Textbox(
@@ -1818,6 +1890,8 @@ with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="MapAnything V2 - 3D重建与
                 with gr.Row():
                     detect_all_btn = gr.Button("🔍 检测所有", size="sm")
                     restore_default_btn = gr.Button("↻ 默认", size="sm")
+                
+                gr.Markdown("📌 **提示**: 启用后会在「分割3D」标签页显示彩色分割模型")
 
     # 示例场景（可折叠）
     with gr.Accordion("🖼️ 示例场景", open=False):
@@ -1873,6 +1947,18 @@ with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="MapAnything V2 - 3D重建与
             "✅ 上传完成，点击「开始重建」进行 3D 处理",
         )
     
+    # 处理摄像头拍照
+    def update_gallery_on_camera(image):
+        if image is None:
+            return None, None, None
+        # 将单张图片包装成列表
+        target_dir, image_paths = handle_uploads([image], 1.0)
+        return (
+            target_dir,
+            image_paths,
+            "✅ 摄像头照片已添加，点击「开始重建」进行 3D 处理",
+        )
+    
     def show_resample_button(files):
         """仅当上传的文件包含视频时显示重新采样按钮"""
         if not files:
@@ -1955,6 +2041,13 @@ with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="MapAnything V2 - 3D重建与
         outputs=[resample_btn],
     )
     
+    # 摄像头拍照事件
+    camera_input.change(
+        fn=update_gallery_on_camera,
+        inputs=[camera_input],
+        outputs=[target_dir_output, image_gallery, log_output]
+    )
+    
     # 滑块改变时显示重新采样按钮（仅当已上传文件时）
     s_time_interval.change(
         fn=show_resample_button,
@@ -2076,12 +2169,16 @@ with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="MapAnything V2 - 3D重建与
     )
 
 # 启动信息
-print("\n" + "="*60)
+print("\n" + "="*70)
 print("🚀 MapAnything V2 - 3D重建与物体分割")
-print("="*60)
+print("="*70)
 print("📊 核心技术: 自适应DBSCAN聚类 + 多视图融合")
 print(f"🔧 质量控制: 置信度≥{MIN_DETECTION_CONFIDENCE} | 面积≥{MIN_MASK_AREA}px")
 print(f"🎯 聚类半径: 沙发{DBSCAN_EPS_CONFIG['sofa']}m | 桌子{DBSCAN_EPS_CONFIG['table']}m | 窗户{DBSCAN_EPS_CONFIG['window']}m | 默认{DBSCAN_EPS_CONFIG['default']}m")
-print("="*60 + "\n")
+print("\n💡 分割配置 (CPU优化):")
+print(f"   - 检测模型: {GROUNDING_DINO_MODEL_ID} (CPU)")
+print(f"   - 分割模型: {SAM_MODEL_ID} (MobileSAM, 10MB, CPU)")
+print(f"   - 运行设备: CPU (不占用GPU资源，适合分离部署)")
+print("="*70 + "\n")
 
 demo.queue(max_size=20).launch(show_error=True, share=True, ssr_mode=False)