models/image_analysis.py

# models/image_analysis.py

from PIL import Image
import numpy as np
from transformers import AutoImageProcessor, AutoModelForImageClassification
from .logging_config import logger

# Initialize real estate classification model
try:
    processor = AutoImageProcessor.from_pretrained("andupets/real-estate-image-classification")
    model = AutoModelForImageClassification.from_pretrained("andupets/real-estate-image-classification")
    has_model = True
    logger.info("Real estate classification model loaded successfully")
except Exception as e:
    logger.error(f"Error loading real estate classification model: {str(e)}")
    has_model = False

def analyze_image(image):
    try:
        if has_model:
            img_rgb = image.convert('RGB')
            inputs = processor(images=img_rgb, return_tensors="pt")
            outputs = model(**inputs)
            logits = outputs.logits
            probs = logits.softmax(dim=1).detach().numpy()[0]
            
            # Get the highest confidence prediction
            max_prob_idx = probs.argmax()
            max_prob = probs[max_prob_idx]
            predicted_label = model.config.id2label[max_prob_idx]
            
            # Check if it's a real estate image (confidence > 0.5)
            is_real_estate = max_prob > 0.5
            
            quality = assess_image_quality(image)
            is_ai_generated = detect_ai_generated_image(image)

            return {
                'is_property_related': is_real_estate,
                'property_confidence': float(max_prob),
                'predicted_label': predicted_label,
                'top_predictions': [
                    {'label': model.config.id2label[i], 'confidence': float(prob)}
                    for i, prob in enumerate(probs)
                ],
                'image_quality': quality,
                'is_ai_generated': is_ai_generated,
                'authenticity_score': 0.95 if not is_ai_generated else 0.60
            }
        else:
            logger.warning("Real estate classification model unavailable")
            return {
                'is_property_related': False,
                'property_confidence': 0.0,
                'predicted_label': 'unknown',
                'top_predictions': [],
                'image_quality': assess_image_quality(image),
                'is_ai_generated': False,
                'authenticity_score': 0.5
            }
    except Exception as e:
        logger.error(f"Error analyzing image: {str(e)}")
        return {
            'is_property_related': False,
            'property_confidence': 0.0,
            'predicted_label': 'error',
            'top_predictions': [],
            'image_quality': {'resolution': 'unknown', 'quality_score': 0},
            'is_ai_generated': False,
            'authenticity_score': 0.0,
            'error': str(e)
        }

def detect_ai_generated_image(image):
    try:
        img_array = np.array(image)
        if len(img_array.shape) == 3:
            gray = np.mean(img_array, axis=2)
        else:
            gray = img_array
        noise = gray - np.mean(gray)
        noise_std = np.std(noise)
        width, height = image.size
        perfect_dimensions = (width % 64 == 0 and height % 64 == 0)
        has_exif = hasattr(image, '_getexif') and image._getexif() is not None
        return noise_std < 0.05 or perfect_dimensions or not has_exif
    except Exception as e:
        logger.error(f"Error detecting AI-generated image: {str(e)}")
        return False

def assess_image_quality(img):
    try:
        width, height = img.size
        resolution = width * height
        quality_score = min(100, resolution // 20000)
        return {
            'resolution': f"{width}x{height}",
            'quality_score': quality_score
        }
    except Exception as e:
        logger.error(f"Error assessing image quality: {str(e)}")
        return {
            'resolution': 'unknown',
            'quality_score': 0
        }