tomatoseveritysegmentationmodel-l5n39e/app/api/routes/model.py

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.orm import Session
from typing import List, Dict, Any
import json
from pathlib import Path

from app.db.session import get_db
from app.db.crud_tomato import tomato_image, analysis_result
from app.schemas.tomato import (
    AnalysisResponse,
    AnalysisResultCreate,
    SeverityDetailCreate
)
from app.services.model import segmentation_model

router = APIRouter()


@router.post("/analyze/{image_id}", response_model=AnalysisResponse)
def analyze_tomato_image(
    image_id: str,
    db: Session = Depends(get_db)
):
    """
    Analyze a tomato image to detect disease severity.
    
    This endpoint processes the image using the segmentation model
    and returns detailed analysis results, including disease severity
    classification and segmentation masks.
    """
    # Get the image from database
    db_image = tomato_image.get(db=db, id=image_id)
    if db_image is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Image not found"
        )
    
    # Check if image file exists
    image_path = db_image.file_path
    if not Path(image_path).exists():
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Image file not found on server"
        )
    
    # Analyze the image
    analysis_results = segmentation_model.analyze_image(image_path)
    
    if not analysis_results.get("success", False):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail=analysis_results.get("error", "Failed to analyze image")
        )
    
    # Store analysis results in database
    analysis_data = {
        "image_id": db_image.id,
        "model_name": analysis_results["model_name"],
        "model_version": analysis_results["model_version"],
        "primary_severity": analysis_results["primary_severity"],
        "severity_confidence": analysis_results["severity_confidence"],
        "segmentation_data": json.dumps(analysis_results["segmentation_data"]),
        "processing_time_ms": analysis_results["processing_time_ms"]
    }
    
    analysis_in = AnalysisResultCreate(**analysis_data)
    
    # Create severity details
    severity_details_data = []
    for detail in analysis_results["severity_details"]:
        severity_details_data.append(
            SeverityDetailCreate(
                severity_class=detail["severity_class"],
                confidence=detail["confidence"],
                affected_area_percentage=detail["affected_area_percentage"],
                analysis_id=""  # Will be set after analysis is created
            )
        )
    
    # Create analysis result with details
    db_analysis = analysis_result.create_with_details(
        db=db, 
        analysis_in=analysis_in, 
        details_in=severity_details_data
    )
    
    # Get fresh data with relationships loaded
    db_analysis = analysis_result.get(db=db, id=db_analysis.id)
    
    # Prepare response
    response = AnalysisResponse(
        id=db_analysis.id,
        image=db_image,
        primary_severity=db_analysis.primary_severity,
        severity_confidence=db_analysis.severity_confidence,
        severity_details=db_analysis.severity_details,
        segmentation_data=db_analysis.segmentation_data,
        processed_at=db_analysis.processed_at,
        model_name=db_analysis.model_name,
        model_version=db_analysis.model_version
    )
    
    return response


@router.get("/results/{image_id}", response_model=List[AnalysisResponse])
def get_analysis_results(
    image_id: str,
    db: Session = Depends(get_db)
):
    """
    Get all analysis results for a specific tomato image.
    """
    # Check if image exists
    db_image = tomato_image.get(db=db, id=image_id)
    if db_image is None:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail="Image not found"
        )
    
    # Get analysis results
    analyses = analysis_result.get_for_image(db=db, image_id=image_id)
    
    # Prepare response
    responses = []
    for db_analysis in analyses:
        responses.append(
            AnalysisResponse(
                id=db_analysis.id,
                image=db_image,
                primary_severity=db_analysis.primary_severity,
                severity_confidence=db_analysis.severity_confidence,
                severity_details=db_analysis.severity_details,
                segmentation_data=db_analysis.segmentation_data,
                processed_at=db_analysis.processed_at,
                model_name=db_analysis.model_name,
                model_version=db_analysis.model_version
            )
        )
    
    return responses


@router.get("/info", response_model=Dict[str, Any])
def get_model_info():
    """
    Get information about the tomato severity segmentation model.
    """
    return {
        "name": segmentation_model.model_name,
        "version": segmentation_model.model_version,
        "description": "Tomato disease severity segmentation model",
        "input_format": {
            "type": "image",
            "size": segmentation_model.input_size,
            "supported_formats": ["JPEG", "PNG"]
        },
        "severity_classes": segmentation_model.severity_classes,
        "capabilities": [
            "disease classification",
            "severity assessment",
            "leaf segmentation"
        ]
    }