chetantiwari/Computer-Vision-Project-Work
1
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 