Classify Industrial Components with Zero-Shot Vision Using CLIP and Detectron2

                      
                     
"""
Production implementation for Classifying Industrial Components with Zero-Shot Vision Using CLIP and Detectron2.
Provides secure, scalable operations.
"""

from typing import Dict, Any, List, Tuple
import os
import logging
import requests
import torch
from torchvision import transforms
from PIL import Image
from detectron2 import model_zoo
from detectron2.config import get_cfg
from detectron2.engine import DefaultPredictor
from detectron2.data import MetadataCatalog

# Setting up logging for tracking execution flow
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class Config:
    """
    Configuration class to manage environment variables.
    """
    def __init__(self):
        self.model_url: str = os.getenv('MODEL_URL', 'COCO-InstanceSegmentation/mask_rcnn_R_101_FPN_3x.yaml')
        self.device: str = os.getenv('DEVICE', 'cuda' if torch.cuda.is_available() else 'cpu')

cfg = Config()  # Initialize configuration

def validate_input(data: Dict[str, Any]) -> bool:
    """Validate input data for classification.
    
    Args:
        data: Input data containing image URL or path.
    Returns:
        True if valid.
    Raises:
        ValueError: If validation fails.
    """
    if 'image_path' not in data:
        raise ValueError('Missing image_path')  # Ensure image path is provided
    return True

def fetch_image(image_path: str) -> Image:
    """Fetch image from local or URL.
    
    Args:
        image_path: Local path or URL for the image.
    Returns:
        PIL Image object.
    Raises:
        RuntimeError: If image cannot be fetched.
    """
    try:
        if image_path.startswith('http'):
            response = requests.get(image_path)
            response.raise_for_status()
            return Image.open(BytesIO(response.content)).convert('RGB')  # Convert to RGB
        else:
            return Image.open(image_path).convert('RGB')  # Local image
    except Exception as e:
        logger.error(f'Error fetching image: {str(e)}')
        raise RuntimeError('Failed to fetch image.')  # Log and raise error

def preprocess_image(image: Image) -> torch.Tensor:
    """Preprocess image for model input.
    
    Args:
        image: The input PIL Image.
    Returns:
        Tensor representation of the image.
    """
    transform = transforms.Compose([
        transforms.Resize((640, 640)),  # Resize to model input size
        transforms.ToTensor(),  # Convert to tensor
    ])
    return transform(image).unsqueeze(0)  # Add batch dimension

def get_predictor() -> DefaultPredictor:
    """Initialize model predictor for CLIP and Detectron2.
    
    Returns:
        DefaultPredictor object.
    """
    cfg = get_cfg()
    cfg.merge_from_file(model_zoo.get_config_file(cfg.model_url))  # Load config
    cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(cfg.model_url)  # Load weights
    cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5  # Set threshold for this model
    cfg.MODEL.DEVICE = cfg.device  # Set device
    return DefaultPredictor(cfg)  # Return predictor

def classify_components(image_path: str) -> List[Dict[str, Any]]:
    """Classify components in the image.
    
    Args:
        image_path: Path to the image to classify.
    Returns:
        List of detected components with their labels and scores.
    Raises:
        RuntimeError: If classification fails.
    """
    validate_input({'image_path': image_path})  # Validate input
    image = fetch_image(image_path)  # Fetch image
    tensor_image = preprocess_image(image)  # Preprocess
    predictor = get_predictor()  # Get model predictor

    outputs = predictor(tensor_image)  # Make prediction
    # Process outputs to extract relevant information
    instances = outputs['instances'].to(torch.device(cfg.device))
    scores = instances.scores.cpu().numpy()  # Get scores
    classes = instances.pred_classes.cpu().numpy()  # Get classes
    return [{'label': MetadataCatalog.get(cfg.model_url).thing_classes[class_id], 'score': score.item()}
            for class_id, score in zip(classes, scores)]  # Return results

def format_output(results: List[Dict[str, Any]]) -> str:
    """Format classification results into a string.
    
    Args:
        results: List of classification results.
    Returns:
        Formatted string for output.
    """
    return '\n'.join([f'Component: {result['label']}, Score: {result['score']:.2f}' for result in results])  # Format results

if __name__ == '__main__':
    try:
        image_path = 'path/to/your/image.jpg'  # Example image path
        results = classify_components(image_path)  # Classify components
        output = format_output(results)  # Format results
        print(output)  # Print results
    except Exception as e:
        logger.error(f'Error in processing: {str(e)}')  # Log error