Team Ai
Apppublic

chetantiwari/Computer-Vision-Project-Work

sourceHugging Faceupdated 1y agoView on Hugging Face
1likes
main.py115 linesDownload Raw Back to app
1import os2os.environ["TRANSFORMERS_CACHE"] = "/tmp/.cache/huggingface/transformers"3os.environ["HF_HOME"] = "/tmp/.cache/huggingface"4os.makedirs("/tmp/.cache/huggingface/transformers", exist_ok=True)5 6 7from fastapi import FastAPI, File, UploadFile, HTTPException8from fastapi.middleware.cors import CORSMiddleware9from fastapi.staticfiles import StaticFiles10from fastapi.responses import HTMLResponse11import uvicorn12from PIL import Image13import io14import asyncio15from typing import Dict, Any16 17from app.models.clothing_detector import ClothingDetector18from app.models.attribute_extractor import AttributeExtractor19from app.models.color_analyzer import ColorAnalyzer20from app.schemas.response import ClothingAnalysisResponse21from app.utils.image_processing import preprocess_image22 23app = FastAPI(title="Clothing Attribute Detection API", version="1.0.0")24 25# Add CORS middleware26app.add_middleware(27    CORSMiddleware,28    allow_origins=["*"],29    allow_credentials=True,30    allow_methods=["*"],31    allow_headers=["*"],32)33 34# Mount static files35app.mount("/static", StaticFiles(directory="frontend"), name="static")36 37# Initialize models (loaded once at startup)38clothing_detector = None39attribute_extractor = None40color_analyzer = None41 42@app.on_event("startup")43async def load_models():44    global clothing_detector, attribute_extractor, color_analyzer45    print("Loading models...")46    47    clothing_detector = ClothingDetector()48    attribute_extractor = AttributeExtractor()49    color_analyzer = ColorAnalyzer()50    51    print("Models loaded successfully!")52 53@app.get("/", response_class=HTMLResponse)54async def read_root():55    with open("frontend/index.html", "r", encoding="utf-8") as f:56        html = f.read()57    return HTMLResponse(html)58 59@app.get("/health")60async def health_check():61    return {"status": "healthy", "message": "Clothing Attribute Detection API is running"}62 63@app.post("/analyze", response_model=ClothingAnalysisResponse)64async def analyze_clothing(file: UploadFile = File(...)):65    try:66        # Validate file type67        if not file.content_type.startswith("image/"):68            raise HTTPException(status_code=400, detail="File must be an image")69        70        # Read and preprocess image71        image_bytes = await file.read()72        image = Image.open(io.BytesIO(image_bytes))73        processed_image = preprocess_image(image)74        75        # Run analysis in parallel76        detection_task = asyncio.create_task(77            clothing_detector.detect_clothing_items(processed_image)78        )79        attribute_task = asyncio.create_task(80            attribute_extractor.extract_attributes(processed_image)81        )82        color_task = asyncio.create_task(83            color_analyzer.analyze_colors(processed_image)84        )85        86        # Wait for all tasks to complete87        clothing_items, attributes, color_analysis = await asyncio.gather(88            detection_task, attribute_task, color_task89        )90        91        # Combine results92        result = {93            "status": "success",94            "clothing_items": clothing_items,95            "style_classification": attributes.get("style", "unknown"),96            "formality": attributes.get("formality", "unknown"), 97            "texture": attributes.get("texture", "unknown"),98            "dominant_colors": color_analysis["dominant_colors"],99            "color_distribution": color_analysis["color_distribution"],100            "detailed_attributes": attributes,101            "confidence_scores": {102                "overall": 0.85,103                "style": attributes.get("confidence", 0.8),104                "color": color_analysis.get("confidence", 0.9)105            }106        }107        108        return ClothingAnalysisResponse(**result)109        110    except Exception as e:111        raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")112 113if __name__ == "__main__":114    uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=True)115