Upload 5 files

main.py

@@ -0,0 +1,346 @@
+"""
+FastAPI application for YOLO PyTorch image classification
+Accepts image uploads and returns verdict (green/red) based on detected classes
+Includes GPU detection and automatic device selection
+"""
+
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import List, Dict, Optional
+import numpy as np
+import cv2
+from io import BytesIO
+from PIL import Image
+import logging
+import base64
+
+from yolo_inference_pytorch import YOLOInference
+from verdict_logic import VerdictEngine
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI(
+    title="YOLO PyTorch Image Classifier - Engineer Selfie",
+    description="Upload an image to get classification verdict (green/red) based on detected objects. Supports GPU acceleration.",
+    version="2.0.0"
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class Base64ImageRequest(BaseModel):
+    image: str
+    
+    class Config:
+        json_schema_extra = {
+            "example": {
+                "image": "data:image/jpeg;base64,/9j/4AAQSkZJRg..."
+            }
+        }
+
+class Detection(BaseModel):
+    class_id: int
+    class_name: str
+    confidence: float
+    bbox: List[float]
+
+class VerdictResponse(BaseModel):
+    verdict: str
+    confidence: float
+    detections: List[Detection]
+    message: str
+    image_size: Dict[str, int]
+
+yolo_model = None
+verdict_engine = None
+
+@app.on_event("startup")
+async def startup_event():
+    """Initialize models on startup"""
+    global yolo_model, verdict_engine
+    
+    try:
+        import os
+        model_path = "model/best.pt"
+        
+        if not os.path.exists(model_path):
+            logger.warning(f"⚠️  YOLO model not found at {model_path}")
+            logger.warning("⚠️  The API will start but /predict endpoint will not work until you add the model")
+            logger.warning("⚠️  Please place your best.pt file in the model/ directory")
+            
+            verdict_engine = VerdictEngine(
+                green_classes=[],
+                red_classes=[],
+                rules_path="config/verdict_rules.json"
+            )
+            logger.info("Verdict engine initialized successfully")
+            return
+        
+        logger.info("Loading YOLO PyTorch model with GPU detection...")
+        yolo_model = YOLOInference(
+            model_path=model_path,
+            labels_path="model/labels.txt",
+            conf_threshold=0.5,
+            iou_threshold=0.45,
+            device=None
+        )
+        logger.info("YOLO model loaded successfully")
+        logger.info(f"Using confidence threshold: 0.5 (50%)")
+        
+        logger.info("Initializing verdict engine...")
+        verdict_engine = VerdictEngine(
+            green_classes=[],
+            red_classes=[],
+            rules_path="config/verdict_rules.json"
+        )
+        logger.info("Verdict engine initialized successfully")
+        
+    except Exception as e:
+        logger.error(f"Error during startup: {str(e)}")
+        raise
+
+@app.get("/")
+async def root():
+    """Root endpoint with API information"""
+    return {
+        "message": "YOLO PyTorch Image Classifier API - Engineer Selfie",
+        "version": "2.0.0",
+        "description": "Upload images to check for object detection and get approval verdicts. GPU-accelerated when available.",
+        "endpoints": {
+            "check_image_base64": "/check-image (POST with base64)",
+            "check_image_upload": "/check-image/upload (POST with file)",
+            "check_images_batch": "/check-images/batch (POST with multiple files)",
+            "health": "/health (GET)",
+            "model_info": "/model/info (GET)",
+            "device_info": "/device/info (GET)"
+        },
+        "usage": {
+            "base64": "POST /check-image with JSON body: {\"image\": \"base64_string\"}",
+            "file_upload": "POST /check-image/upload with multipart/form-data",
+            "batch": "POST /check-images/batch with multiple files"
+        }
+    }
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "model_loaded": yolo_model is not None,
+        "verdict_engine_loaded": verdict_engine is not None,
+        "device": yolo_model.device if yolo_model else "not loaded"
+    }
+
+@app.get("/device/info")
+async def device_info():
+    """Get GPU/device information"""
+    if yolo_model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    return yolo_model.get_device_info()
+
+@app.get("/model/info")
+async def model_info():
+    """Get model information"""
+    if yolo_model is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    return {
+        "model_type": "YOLO PyTorch",
+        "model_path": yolo_model.model_path,
+        "classes": yolo_model.get_class_names(),
+        "num_classes": len(yolo_model.get_class_names()),
+        "input_size": yolo_model.get_input_size(),
+        "confidence_threshold": yolo_model.conf_threshold,
+        "iou_threshold": yolo_model.iou_threshold,
+        "device_info": yolo_model.get_device_info()
+    }
+
+def process_image_to_numpy(image_data: bytes) -> tuple[np.ndarray, int, int]:
+    """Convert image bytes to numpy array for processing"""
+    image = Image.open(BytesIO(image_data))
+    
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
+    
+    image_np = np.array(image)
+    
+    if len(image_np.shape) == 3 and image_np.shape[2] == 3:
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_RGB2BGR)
+    elif len(image_np.shape) == 2:
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_GRAY2BGR)
+    
+    original_height, original_width = image_np.shape[:2]
+    return image_np, original_width, original_height
+
+
+def create_verdict_response(detections: List[Dict], verdict_result: Dict, width: int, height: int) -> VerdictResponse:
+    """Create standardized verdict response"""
+    formatted_detections = [
+        Detection(
+            class_id=det["class_id"],
+            class_name=det["class_name"],
+            confidence=float(det["confidence"]),
+            bbox=[float(x) for x in det["bbox"]]
+        )
+        for det in detections
+    ]
+    
+    return VerdictResponse(
+        verdict=verdict_result["verdict"],
+        confidence=verdict_result["confidence"],
+        detections=formatted_detections,
+        message=verdict_result["message"],
+        image_size={
+            "width": width,
+            "height": height
+        }
+    )
+
+
+@app.post("/check-image", response_model=VerdictResponse)
+async def check_image_base64(request: Base64ImageRequest):
+    """
+    CHECK IMAGE (Base64 Format)
+    
+    Send a base64 encoded image and get a verdict (GREEN/RED) based on detected objects.
+    """
+    if yolo_model is None or verdict_engine is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    try:
+        base64_string = request.image
+        
+        if "base64," in base64_string:
+            base64_string = base64_string.split("base64,")[1]
+        
+        image_bytes = base64.b64decode(base64_string)
+        logger.info("Processing base64 encoded image")
+        
+        image_np, original_width, original_height = process_image_to_numpy(image_bytes)
+        logger.info(f"Image size: {original_width}x{original_height}")
+        
+        detections = yolo_model.predict(image_np)
+        logger.info(f"Found {len(detections)} detections")
+        
+        verdict_result = verdict_engine.get_verdict(detections)
+        
+        response = create_verdict_response(detections, verdict_result, original_width, original_height)
+        
+        logger.info(f"Verdict: {verdict_result['verdict']} (confidence: {verdict_result['confidence']:.2f})")
+        return response
+        
+    except base64.binascii.Error:
+        raise HTTPException(status_code=400, detail="Invalid base64 image data")
+    except Exception as e:
+        logger.error(f"Error processing image: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
+
+
+@app.post("/check-image/upload", response_model=VerdictResponse)
+async def check_image_upload(file: UploadFile = File(...)):
+    """
+    CHECK IMAGE (File Upload)
+    
+    Upload an image file and get a verdict (GREEN/RED) based on detected objects.
+    """
+    if yolo_model is None or verdict_engine is None:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
+    if not file.content_type.startswith("image/"):
+        raise HTTPException(
+            status_code=400,
+            detail=f"Invalid file type: {file.content_type}. Please upload an image."
+        )
+    
+    try:
+        logger.info(f"Processing image: {file.filename}")
+        contents = await file.read()
+        
+        image_np, original_width, original_height = process_image_to_numpy(contents)
+        logger.info(f"Image size: {original_width}x{original_height}")
+        
+        detections = yolo_model.predict(image_np)
+        logger.info(f"Found {len(detections)} detections")
+        
+        verdict_result = verdict_engine.get_verdict(detections)
+        
+        response = create_verdict_response(detections, verdict_result, original_width, original_height)
+        
+        logger.info(f"Verdict: {verdict_result['verdict']} (confidence: {verdict_result['confidence']:.2f})")
+        return response
+        
+    except Exception as e:
+        logger.error(f"Error processing image: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
+
+
+@app.post("/predict", response_model=VerdictResponse)
+async def predict(file: UploadFile = File(...)):
+    """
+    [LEGACY] Upload an image and get classification verdict
+    
+    Note: Use /check-image/upload instead for clearer API naming
+    """
+    return await check_image_upload(file)
+
+@app.post("/check-images/batch")
+async def check_images_batch(files: List[UploadFile] = File(...)):
+    """
+    CHECK MULTIPLE IMAGES (Batch Upload)
+    
+    Upload multiple image files (max 10) and get verdicts for each.
+    """
+    if len(files) > 10:
+        raise HTTPException(
+            status_code=400,
+            detail="Maximum 10 images allowed per batch"
+        )
+    
+    results = []
+    for file in files:
+        try:
+            result = await check_image_upload(file)
+            results.append({
+                "filename": file.filename,
+                "result": result
+            })
+        except Exception as e:
+            results.append({
+                "filename": file.filename,
+                "error": str(e)
+            })
+    
+    return {"results": results}
+
+
+@app.post("/predict/batch")
+async def predict_batch(files: List[UploadFile] = File(...)):
+    """
+    [LEGACY] Upload multiple images and get classification verdicts
+    
+    Note: Use /check-images/batch instead for clearer API naming
+    """
+    return await check_images_batch(files)
+
+if __name__ == "__main__":
+    import uvicorn
+    import sys
+    
+    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8000
+    
+    print(f"\n{'='*70}")
+    print(f"🚀 Starting Engineer Selfie API on http://0.0.0.0:{port}")
+    print(f"📖 API docs available at http://localhost:{port}/docs")
+    print(f"🖥️  GPU acceleration will be used if available")
+    print(f"{'='*70}\n")
+    
+    uvicorn.run(app, host="0.0.0.0", port=port)

model/best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0e9194db6fc94790b6902f48e8a5e9c04e025b47ca745bd0b3569d3814a1509c
+size 6223722

model/labels.txt

@@ -0,0 +1,3 @@
+jio jersey
+jio logo
+person

requirements.txt

@@ -0,0 +1,18 @@
+# FastAPI and server dependencies
+fastapi==0.109.0
+uvicorn[standard]==0.27.0
+python-multipart==0.0.6
+
+# Machine Learning and Image Processing
+opencv-python>=4.10.0
+numpy>=1.24.0
+Pillow==10.2.0
+torch>=2.0.0
+torchvision>=0.15.0
+ultralytics>=8.0.0
+
+# Data validation and utilities
+pydantic==2.5.3
+
+# Logging and monitoring (optional)
+python-json-logger==2.0.7

verdict_logic.py

@@ -0,0 +1,327 @@
+"""
+Verdict Logic Engine
+Determines green/red verdict based on detected classes
+"""
+
+import json
+from typing import List, Dict, Optional
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class VerdictEngine:
+    """
+    Engine to determine verdict (green/red) based on detected objects
+    """
+    
+    def __init__(
+        self,
+        green_classes: Optional[List[str]] = None,
+        red_classes: Optional[List[str]] = None,
+        rules_path: Optional[str] = None
+    ):
+        """
+        Initialize verdict engine
+        
+        Args:
+            green_classes: List of class names that trigger green verdict
+            red_classes: List of class names that trigger red verdict
+            rules_path: Optional path to JSON file with custom rules
+        """
+        self.green_classes = green_classes or []
+        self.red_classes = red_classes or []
+        self.custom_rules = {}
+        
+        # Load custom rules if provided
+        if rules_path:
+            try:
+                self.load_rules(rules_path)
+            except FileNotFoundError:
+                logger.warning(f"Rules file not found: {rules_path}")
+        
+        # Default configuration if no rules specified
+        if not self.green_classes and not self.red_classes and not self.custom_rules:
+            logger.info("No verdict rules specified. Using default configuration.")
+            self._setup_default_rules()
+    
+    def _setup_default_rules(self):
+        """
+        Setup default verdict rules
+        Customize this based on your specific use case
+        """
+        # Example default rules:
+        # - If all 3 classes detected: GREEN
+        # - If only specific classes detected: RED
+        # - If no classes detected: RED
+        
+        self.custom_rules = {
+            "all_classes_required": False,  # Set to True if all 3 classes must be present for GREEN
+            "min_detections_for_green": 1,  # Minimum detections needed for GREEN
+            "min_confidence_for_green": 0.5,  # Minimum confidence for GREEN verdict
+            "priority": "red"  # If both green and red classes detected, which takes priority
+        }
+        
+        logger.info("Default verdict rules configured")
+    
+    def load_rules(self, rules_path: str):
+        """Load verdict rules from JSON file"""
+        with open(rules_path, 'r') as f:
+            rules = json.load(f)
+        
+        self.green_classes = rules.get("green_classes", self.green_classes)
+        self.red_classes = rules.get("red_classes", self.red_classes)
+        self.custom_rules = rules.get("custom_rules", self.custom_rules)
+        
+        logger.info(f"Loaded rules from {rules_path}")
+        logger.info(f"Green classes: {self.green_classes}")
+        logger.info(f"Red classes: {self.red_classes}")
+    
+    def get_verdict(self, detections: List[Dict]) -> Dict:
+        """
+        Determine verdict based on detections
+        
+        New Logic:
+        - GREEN: All 3 labels present (jio jersey, jio logo, person) AND exactly 1 person
+        - RED: Any condition not met
+        
+        Args:
+            detections: List of detection dictionaries from YOLO
+        
+        Returns:
+            Dictionary with verdict, confidence, and message
+        """
+        if len(detections) == 0:
+            return {
+                "verdict": "red",
+                "confidence": 1.0,
+                "message": "Image rejected. No objects detected in the image. Please ensure the image contains: Jio jersey, Jio logo, and exactly one person."
+            }
+        
+        # Extract detected class names and confidences
+        detected_classes = [det["class_name"] for det in detections]
+        confidences = [det["confidence"] for det in detections]
+        avg_confidence = sum(confidences) / len(confidences)
+        
+        # Check which required labels are present
+        unique_classes = set(detected_classes)
+        
+        # Count number of persons detected
+        person_count = detected_classes.count("person")
+        
+        # Check conditions
+        has_jio_jersey = "jio jersey" in unique_classes
+        has_jio_logo = "jio logo" in unique_classes
+        has_person = "person" in unique_classes
+        has_exactly_one_person = person_count == 1
+        
+        # All conditions must be satisfied for GREEN
+        all_labels_present = has_jio_jersey and has_jio_logo and has_person
+        all_conditions_met = all_labels_present and has_exactly_one_person
+        
+        logger.info(f"Detected classes: {detected_classes}")
+        logger.info(f"Unique classes: {unique_classes}")
+        logger.info(f"Person count: {person_count}")
+        logger.info(f"All conditions met: {all_conditions_met}")
+        
+        # Build user-friendly message
+        if all_conditions_met:
+            message = "Image approved! All requirements met: Jio jersey found, Jio logo found, and exactly one person detected."
+            verdict = "green"
+        else:
+            # Build detailed explanation of what's missing or wrong
+            issues = []
+            
+            if not has_jio_jersey:
+                issues.append("Jio jersey not detected")
+            if not has_jio_logo:
+                issues.append("Jio logo not detected")
+            if not has_person:
+                issues.append("No person detected")
+            elif person_count == 0:
+                issues.append("No person detected")
+            elif person_count > 1:
+                issues.append(f"Multiple people detected ({person_count} found, need exactly 1)")
+            
+            # Build the message
+            if issues:
+                message = f"Image rejected. {'. '.join(issues)}. Requirements: Jio jersey, Jio logo, and exactly one person must be present."
+            else:
+                message = "Image rejected. Please ensure all requirements are met: Jio jersey, Jio logo, and exactly one person."
+            
+            verdict = "red"
+        
+        return {
+            "verdict": verdict,
+            "confidence": avg_confidence,
+            "message": message
+        }
+    
+    def _apply_verdict_logic(
+        self,
+        detected_classes: List[str],
+        unique_classes: set,
+        num_unique_classes: int,
+        max_confidence: float,
+        avg_confidence: float,
+        num_detections: int
+    ) -> tuple:
+        """
+        Apply verdict logic based on detected classes
+        
+        Returns:
+            Tuple of (verdict, message)
+        """
+        # Strategy 1: Check if specific green/red classes are defined
+        if self.green_classes or self.red_classes:
+            return self._apply_class_based_logic(
+                detected_classes, unique_classes, avg_confidence
+            )
+        
+        # Strategy 2: Use custom rules
+        if self.custom_rules:
+            return self._apply_custom_rules(
+                num_unique_classes, num_detections, avg_confidence
+            )
+        
+        # Strategy 3: Default logic (customize based on your needs)
+        return self._apply_default_logic(
+            num_unique_classes, num_detections, avg_confidence
+        )
+    
+    def _apply_class_based_logic(
+        self,
+        detected_classes: List[str],
+        unique_classes: set,
+        avg_confidence: float
+    ) -> tuple:
+        """Apply logic based on specific green/red class lists"""
+        
+        # Check for red classes first (if priority is red)
+        priority = self.custom_rules.get("priority", "red")
+        
+        if priority == "red" and self.red_classes:
+            red_detected = any(cls in self.red_classes for cls in detected_classes)
+            if red_detected:
+                return (
+                    "red",
+                    f"Detected prohibited object(s): {', '.join(unique_classes & set(self.red_classes))}"
+                )
+        
+        # Check for green classes
+        if self.green_classes:
+            green_detected = any(cls in self.green_classes for cls in detected_classes)
+            if green_detected:
+                return (
+                    "green",
+                    f"Detected approved object(s): {', '.join(unique_classes & set(self.green_classes))}"
+                )
+        
+        # Check for red classes (if priority is green)
+        if priority == "green" and self.red_classes:
+            red_detected = any(cls in self.red_classes for cls in detected_classes)
+            if red_detected:
+                return (
+                    "red",
+                    f"Detected prohibited object(s): {', '.join(unique_classes & set(self.red_classes))}"
+                )
+        
+        # Default to red if no matching classes
+        return (
+            "red",
+            f"Detected objects do not match approved criteria: {', '.join(unique_classes)}"
+        )
+    
+    def _apply_custom_rules(
+        self,
+        num_unique_classes: int,
+        num_detections: int,
+        avg_confidence: float
+    ) -> tuple:
+        """Apply custom rules from configuration"""
+        
+        all_classes_required = self.custom_rules.get("all_classes_required", False)
+        min_detections = self.custom_rules.get("min_detections_for_green", 1)
+        min_confidence = self.custom_rules.get("min_confidence_for_green", 0.5)
+        
+        # Check if all 3 classes are required
+        if all_classes_required:
+            if num_unique_classes == 3:
+                if avg_confidence >= min_confidence:
+                    return (
+                        "green",
+                        f"All required classes detected with confidence {avg_confidence:.2%}"
+                    )
+                else:
+                    return (
+                        "red",
+                        f"All classes detected but confidence too low: {avg_confidence:.2%}"
+                    )
+            else:
+                return (
+                    "red",
+                    f"Only {num_unique_classes}/3 required classes detected"
+                )
+        
+        # Check minimum detections
+        if num_detections >= min_detections and avg_confidence >= min_confidence:
+            return (
+                "green",
+                f"Detected {num_detections} object(s) with confidence {avg_confidence:.2%}"
+            )
+        
+        return (
+            "red",
+            f"Insufficient detections ({num_detections}) or confidence ({avg_confidence:.2%})"
+        )
+    
+    def _apply_default_logic(
+        self,
+        num_unique_classes: int,
+        num_detections: int,
+        avg_confidence: float
+    ) -> tuple:
+        """
+        Apply default verdict logic
+        
+        Default strategy:
+        - GREEN: If at least 1 class detected with good confidence
+        - RED: If no detections or low confidence
+        
+        Customize this method based on your specific requirements
+        """
+        
+        # Example 1: Simple presence-based logic
+        if num_detections >= 1 and avg_confidence >= 0.5:
+            return (
+                "green",
+                f"Detected {num_detections} object(s) across {num_unique_classes} class(es)"
+            )
+        
+        # Example 2: All classes must be present
+        # if num_unique_classes == 3 and avg_confidence >= 0.6:
+        #     return ("green", "All 3 classes detected with high confidence")
+        
+        # Example 3: Specific number of detections
+        # if num_detections >= 2 and avg_confidence >= 0.5:
+        #     return ("green", f"Multiple objects detected ({num_detections})")
+        
+        return (
+            "red",
+            f"Criteria not met: {num_detections} detection(s), {num_unique_classes} class(es)"
+        )
+    
+    def set_green_classes(self, classes: List[str]):
+        """Set classes that trigger green verdict"""
+        self.green_classes = classes
+        logger.info(f"Green classes set to: {classes}")
+    
+    def set_red_classes(self, classes: List[str]):
+        """Set classes that trigger red verdict"""
+        self.red_classes = classes
+        logger.info(f"Red classes set to: {classes}")
+    
+    def update_rules(self, rules: Dict):
+        """Update custom rules"""
+        self.custom_rules.update(rules)
+        logger.info(f"Custom rules updated: {self.custom_rules}")