请求将基于Streamlit的衣物检测与描述Web应用转换为FastAPI接口用于部署

I've refactored your code into a production-ready FastAPI backend that handles image uploads, detects clothing items, generates descriptions, and returns structured results. Here's the complete solution:

Full FastAPI Code #

import os
import cv2
import base64
import io
import numpy as np
from PIL import Image
from fastapi import FastAPI, File, UploadFile, HTTPException
from ultralytics import YOLO
from transformers import Blip2ForConditionalGeneration, AutoProcessor
from peft import PeftModel, PeftConfig

# Initialize FastAPI app
app = FastAPI(title="Clothing Detection & Description API", version="1.0")

# Configuration Constants
PREPROCESS_CKP = "./Salesforce/blip2-opt-2.7b"  # Local path or Hugging Face repo ID
PEFT_MODEL_ID = "ManhTien22/blip2-opt-2.7b-Kimono"
YOLO_MODEL_PATH = "./model/last.pt"

# Load models once on app startup (efficient for repeated requests)
@app.on_event("startup")
def load_ml_models():
    """Load detection and captioning models when the API starts"""
    try:
        app.state.yolo_model = YOLO(YOLO_MODEL_PATH)
        print("✅ YOLO detection model loaded successfully")
    except Exception as e:
        print(f"❌ Failed to load YOLO model: {str(e)}")
        raise e  # Prevent app from starting if critical model fails

    try:
        # Load BLIP-2 processor and fine-tuned model
        app.state.blip_processor = AutoProcessor.from_pretrained(PREPROCESS_CKP)
        peft_config = PeftConfig.from_pretrained(PEFT_MODEL_ID)
        base_model = Blip2ForConditionalGeneration.from_pretrained(peft_config.base_model_name_or_path)
        app.state.blip_model = PeftModel.from_pretrained(base_model, PEFT_MODEL_ID)
        print("✅ BLIP-2 captioning model loaded successfully")
    except Exception as e:
        print(f"❌ Failed to load BLIP-2 model: {str(e)}")
        raise e

def detect_and_crop_objects(image_np: np.ndarray, model, confidence: float = 0.4):
    """Detect clothing items and return cropped regions + visualization"""
    results = model.predict(source=image_np, conf=confidence, save=False)
    result = results[0]
    
    cropped_images = []
    cropped_labels = []
    original_img = result.orig_img

    for box in result.boxes:
        # Extract bounding box coordinates
        x1, y1, x2, y2 = map(int, box.xyxy.cpu().numpy()[0])
        cropped = original_img[y1:y2, x1:x2]
        cropped_images.append(cropped)
        
        # Get label for detected object
        label = result.names[int(box.cls.cpu().numpy()[0])]
        cropped_labels.append(label)

    # Return cropped items, labels, and visualization (RGB format)
    return cropped_images, cropped_labels, result.plot()[:, :, ::-1]

def generate_clothing_caption(image: Image.Image, processor, model, label: str = None):
    """Generate descriptive caption for a clothing item using BLIP-2"""
    # Create prompt focused on detected label
    if label:
        prompt = f"What are the key features of this {label} clothing item?"
    else:
        prompt = "Describe the key features of the clothing item in this image."

    # Process input and generate caption
    inputs = processor(images=image, text=prompt, return_tensors="pt")
    generated_ids = model.generate(
        pixel_values=inputs.pixel_values,
        max_length=60,
        num_beams=4,
        early_stopping=True,
        repetition_penalty=1.8
    )
    return processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

@app.get("/", summary="API Health Check")
async def health_check():
    """Verify the API is running and models are loaded"""
    return {
        "status": "healthy",
        "message": "Clothing Detection API is ready to process requests"
    }

@app.post("/predict", summary="Process Image for Clothing Detection & Description")
async def process_image(
    file: UploadFile = File(..., description="Upload JPG/JPEG/PNG image containing clothing"),
    confidence: float = 0.4  # Adjust detection sensitivity (0.0 to 1.0)
):
    # Validate file type
    allowed_types = ["image/jpeg", "image/png", "image/jpg"]
    if file.content_type not in allowed_types:
        raise HTTPException(
            status_code=400,
            detail=f"Invalid file type. Allowed formats: {', '.join(allowed_types)}"
        )

    try:
        # Read and convert uploaded image to numpy array
        image_data = await file.read()
        pil_img = Image.open(io.BytesIO(image_data))
        img_np = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Failed to process image: {str(e)}")

    try:
        # Step 1: Detect and crop clothing items
        cropped_items, labels, detection_visual = detect_and_crop_objects(
            img_np, app.state.yolo_model, confidence=confidence
        )

        if not cropped_items:
            return {
                "message": "No clothing items detected in the image",
                "detection_visual": None,
                "descriptions": []
            }

        # Step 2: Generate descriptions for each cropped item
        descriptions = []
        for crop, label in zip(cropped_items, labels):
            # Convert crop to PIL image for BLIP-2
            crop_pil = Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
            
            # Generate caption
            caption = generate_clothing_caption(
                crop_pil, app.state.blip_processor, app.state.blip_model, label=label
            )

            # Encode cropped image to base64 for frontend display
            buffer = io.BytesIO()
            crop_pil.save(buffer, format="PNG")
            crop_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")

            descriptions.append({
                "label": label.capitalize(),
                "description": caption,
                "cropped_image": f"data:image/png;base64,{crop_base64}"
            })

        # Encode detection visualization to base64
        visual_pil = Image.fromarray(detection_visual)
        buffer = io.BytesIO()
        visual_pil.save(buffer, format="PNG")
        visual_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")

        # Return structured results
        return {
            "message": "Image processed successfully",
            "detection_visual": f"data:image/png;base64,{visual_base64}",
            "descriptions": descriptions
        }

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Processing error: {str(e)}")

Key Changes & Improvements #

1. Removed Streamlit UI Code: All Streamlit-specific components (like st.title, st.file_uploader) were removed since FastAPI is a backend-only framework.
2. Efficient Model Loading: Uses FastAPI's startup event to load models once when the API starts, avoiding repeated model loading for each request.
3. In-Memory Processing: Avoids saving files to disk by processing images directly in memory, improving speed and reducing clutter.
4. Structured API Endpoints:
   • /: Health check endpoint to verify API status
   • /predict: Main endpoint that accepts image uploads and returns detection results + descriptions
5. Error Handling: Proper HTTP status codes and meaningful error messages for common issues (invalid file types, processing failures).
6. Base64 Image Encoding: Returns images as base64 strings so frontend applications can easily display detection visualizations and cropped items.

How to Use #

1. Install Dependencies:

   pip install fastapi uvicorn ultralytics transformers peft pillow opencv-python numpy
   

2. Run the API:

   uvicorn main:app --reload --host 0.0.0.0 --port 8000
   

3. Test the API:
   • Visit http://localhost:8000/docs to use the interactive Swagger UI
   • Upload an image and adjust the confidence threshold to see results
   • Or send a POST request to http://localhost:8000/predict using tools like Postman or curl:

     curl -X POST "http://localhost:8000/predict?confidence=0.4" -F "file=@your_image.jpg"
     

Frontend Integration Tips #

• The API returns JSON with:
  • detection_visual: Base64 string of the image with bounding boxes
  • descriptions: List of objects with labels, descriptions, and cropped images
• Frontend can render base64 images directly using HTML img tags:

  <img src="data:image/png;base64,{base64_string}" alt="Detection Result">
  

备注：内容来源于stack exchange，提问作者Kazuto Kirigaya