app/services/weight_calculation_service.py

# 檔案路徑: app/services/weight_calculation_service.py

import logging
import numpy as np
from PIL import Image
import io
from typing import Dict, Any, List, Optional, Tuple
import torch
from transformers import SamModel, SamProcessor, pipeline

# 設置日誌
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 食物密度表 (g/cm³) - 常見食物的平均密度
FOOD_DENSITY_TABLE = {
    "rice": 0.8,           # 米飯
    "fried_rice": 0.7,     # 炒飯
    "noodles": 0.6,        # 麵條
    "bread": 0.3,          # 麵包
    "meat": 1.0,           # 肉類
    "fish": 1.1,           # 魚類
    "vegetables": 0.4,     # 蔬菜
    "fruits": 0.8,         # 水果
    "soup": 1.0,           # 湯類
    "sushi": 0.9,          # 壽司
    "pizza": 0.6,          # 披薩
    "hamburger": 0.7,      # 漢堡
    "salad": 0.3,          # 沙拉
    "default": 0.8         # 預設密度
}

class WeightCalculationService:
    def __init__(self):
        """初始化重量計算服務"""
        self.sam_model = None
        self.sam_processor = None
        self.dpt_model = None
        self._load_models()
    
    def _load_models(self):
        """載入 SAM 和 DPT 模型"""
        try:
            # 載入 SAM 分割模型
            logger.info("正在載入 SAM 分割模型...")
            self.sam_model = SamModel.from_pretrained("facebook/sam-vit-base")
            self.sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
            
            # 載入 DPT 深度估計模型
            logger.info("正在載入 DPT 深度估計模型...")
            self.dpt_model = pipeline("depth-estimation", model="Intel/dpt-large")
            
            logger.info("SAM 和 DPT 模型載入完成！")
            
        except Exception as e:
            logger.error(f"模型載入失敗: {str(e)}")
            raise
    
    def segment_food_area(self, image: Image.Image, bbox: List[float]) -> np.ndarray:
        """
        使用 SAM 分割食物區域
        
        Args:
            image: PIL Image 物件
            bbox: 邊界框 [x1, y1, x2, y2]
            
        Returns:
            np.ndarray: 食物區域的遮罩
        """
        try:
            # 使用 SAM 進行分割
            inputs = self.sam_processor(image, input_boxes=[bbox], return_tensors="pt")
            
            with torch.no_grad():
                outputs = self.sam_model(**inputs)
            
            # 取得分割遮罩
            masks_tensor = self.sam_processor.image_processor.post_process_masks(
                outputs.pred_masks.sigmoid(), 
                inputs["original_sizes"], 
                inputs["reshaped_input_sizes"]
            )[0]
            
            # 選擇最大的遮罩
            mask = masks_tensor[0].squeeze().cpu().numpy().astype(bool)
            
            logger.info(f"SAM 分割完成，遮罩大小: {mask.shape}")
            return mask
            
        except Exception as e:
            logger.error(f"SAM 分割失敗: {str(e)}")
            # 回傳預設遮罩
            return np.ones((image.height, image.width), dtype=bool)
    
    def estimate_depth(self, image: Image.Image) -> np.ndarray:
        """
        使用 DPT 進行深度估計
        
        Args:
            image: PIL Image 物件
            
        Returns:
            np.ndarray: 深度圖
        """
        try:
            # 使用 DPT 進行深度估計
            depth_result = self.dpt_model(image)
            depth_map = depth_result["depth"]
            
            logger.info(f"DPT 深度估計完成，深度圖大小: {depth_map.shape}")
            return np.array(depth_map)
            
        except Exception as e:
            logger.error(f"DPT 深度估計失敗: {str(e)}")
            # 回傳預設深度圖
            return np.ones((image.height, image.width))
    
    def calculate_volume_and_weight(self, 
                                  mask: np.ndarray, 
                                  depth_map: np.ndarray,
                                  food_name: str,
                                  pixel_ratio: float) -> Tuple[float, float, float]:
        """
        計算體積和重量
        
        Args:
            mask: 食物區域遮罩
            depth_map: 深度圖
            food_name: 食物名稱
            pixel_ratio: 像素比例 (cm/pixel)
            
        Returns:
            Tuple[float, float, float]: (重量, 信心度, 誤差範圍)
        """
        try:
            # 計算食物區域的像素數量
            food_pixels = np.sum(mask)
            logger.info(f"重量計算開始 - 食物: {food_name}, 像素數量: {food_pixels}")
            
            # 計算食物區域的平均深度
            food_depth_values = depth_map[mask]
            if len(food_depth_values) > 0:
                food_depth = np.mean(food_depth_values)
                depth_variance = np.var(food_depth_values)
            else:
                food_depth = 1.0
                depth_variance = 0.0
            
            logger.info(f"深度分析 - 平均深度: {food_depth:.4f}, 深度變異: {depth_variance:.4f}")
            
            # 計算實際面積 (cm²)
            area_cm2 = food_pixels * (pixel_ratio ** 2)
            logger.info(f"面積計算 - 像素比例: {pixel_ratio:.4f}, 實際面積: {area_cm2:.2f} cm²")
            
            # 動態調整形狀因子 (基於深度資訊)
            if depth_variance > 0:
                # 深度變異大，表示食物較立體
                shape_factor = np.clip(0.6 + (depth_variance * 0.2), 0.3, 0.8)
            else:
                # 深度變異小，表示食物較扁平
                shape_factor = np.clip(0.4 + (food_depth * 0.2), 0.2, 0.7)
            
            logger.info(f"形狀因子 - 動態調整: {shape_factor:.4f}")
            
            # 計算體積 (cm³)
            volume_cm3 = shape_factor * (area_cm2 ** 1.5)
            logger.info(f"體積計算 - 估算體積: {volume_cm3:.2f} cm³")
            
            # 取得食物密度
            density = self._get_food_density(food_name)
            logger.info(f"密度查詢 - 食物: {food_name}, 密度: {density} g/cm³")
            
            # 計算重量 (g)
            weight = volume_cm3 * density
            logger.info(f"重量計算 - 原始重量: {weight:.2f} g")
            
            # 合理性檢查和調整
            if weight > 2000:  # 超過 2kg
                logger.warning(f"重量 {weight:.2f}g 過高，進行調整")
                weight = min(weight, 2000)
            elif weight < 10:  # 少於 10g
                logger.warning(f"重量 {weight:.2f}g 過低，進行調整")
                weight = max(weight, 10)
            
            # 計算信心度和誤差範圍
            confidence = self._calculate_confidence(pixel_ratio, depth_variance, food_pixels)
            error_range = self._calculate_error_range(confidence)
            
            logger.info(f"最終結果 - 重量: {weight:.2f}g, 信心度: {confidence:.2f}, 誤差範圍: ±{error_range*100:.1f}%")
            
            return weight, confidence, error_range
            
        except Exception as e:
            logger.error(f"重量計算失敗: {str(e)}")
            return 150.0, 0.3, 0.5  # 預設值
    
    def _get_food_density(self, food_name: str) -> float:
        """根據食物名稱取得密度"""
        food_name_lower = food_name.lower()
        
        # 關鍵字匹配
        for keyword, density in FOOD_DENSITY_TABLE.items():
            if keyword in food_name_lower:
                return density
        
        return FOOD_DENSITY_TABLE["default"]
    
    def _calculate_confidence(self, pixel_ratio: float, depth_variance: float, food_pixels: int) -> float:
        """計算信心度"""
        # 基礎信心度
        base_confidence = 0.6
        
        # 像素比例影響 (比例越合理，信心度越高)
        if 0.005 <= pixel_ratio <= 0.05:
            base_confidence += 0.2
        elif 0.001 <= pixel_ratio <= 0.1:
            base_confidence += 0.1
        
        # 深度變異影響 (適中的變異表示好的深度估計)
        if 0.01 <= depth_variance <= 0.1:
            base_confidence += 0.1
        elif depth_variance > 0.5:
            base_confidence -= 0.1
        
        # 像素數量影響 (適中的像素數量表示好的分割)
        if 1000 <= food_pixels <= 100000:
            base_confidence += 0.1
        elif food_pixels < 100:
            base_confidence -= 0.2
        
        return np.clip(base_confidence, 0.1, 0.95)
    
    def _calculate_error_range(self, confidence: float) -> float:
        """根據信心度計算誤差範圍"""
        # 信心度越高，誤差範圍越小
        if confidence >= 0.8:
            return 0.1  # ±10%
        elif confidence >= 0.6:
            return 0.2  # ±20%
        elif confidence >= 0.4:
            return 0.3  # ±30%
        else:
            return 0.5  # ±50%

# 全域服務實例
weight_calculation_service = WeightCalculationService()

def calculate_food_weight(image_bytes: bytes, 
                         food_name: str, 
                         pixel_ratio: float,
                         bbox: Optional[List[float]] = None) -> Dict[str, Any]:
    """
    計算食物重量的主函數
    
    Args:
        image_bytes: 圖片二進位數據
        food_name: 食物名稱
        pixel_ratio: 像素比例
        bbox: 可選的邊界框，如果沒有提供則使用整個圖片
        
    Returns:
        Dict: 包含重量計算結果的字典
    """
    try:
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        
        # 如果沒有提供邊界框，使用整個圖片
        if bbox is None:
            bbox = [0, 0, image.width, image.height]
        
        logger.info(f"開始計算 {food_name} 的重量，使用邊界框: {bbox}")
        
        # 1. 使用 SAM 分割食物區域
        mask = weight_calculation_service.segment_food_area(image, bbox)
        
        # 2. 使用 DPT 進行深度估計
        depth_map = weight_calculation_service.estimate_depth(image)
        
        # 3. 計算體積和重量
        weight, confidence, error_range = weight_calculation_service.calculate_volume_and_weight(
            mask, depth_map, food_name, pixel_ratio
        )
        
        # 4. 計算誤差範圍
        weight_min = weight * (1 - error_range)
        weight_max = weight * (1 + error_range)
        
        return {
            "food_name": food_name,
            "estimated_weight": round(weight, 1),
            "weight_confidence": confidence,
            "weight_error_range": [round(weight_min, 1), round(weight_max, 1)],
            "pixel_ratio": pixel_ratio,
            "calculation_method": "SAM+DPT",
            "success": True
        }
        
    except Exception as e:
        logger.error(f"重量計算主流程失敗: {str(e)}")
        return {
            "food_name": food_name,
            "estimated_weight": 150.0,
            "weight_confidence": 0.3,
            "weight_error_range": [100.0, 200.0],
            "pixel_ratio": pixel_ratio,
            "calculation_method": "SAM+DPT",
            "success": False,
            "error": str(e)
        }