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 image is None:
            logger.error("No image provided to analyze_image.")
            return {
                'is_property_related': False,
                'property_confidence': 0.0,
                'predicted_label': 'no_image',
                'top_predictions': [],
                'image_quality': {'resolution': 'unknown', 'quality_score': 0},
                'is_ai_generated': False,
                'authenticity_score': 0.0,
                'error': 'No image provided'
            }
        if has_model:
            try:
                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]
                max_prob_idx = probs.argmax()
                max_prob = probs[max_prob_idx]
                predicted_label = model.config.id2label[max_prob_idx]
                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
                }
            except Exception as e:
                logger.error(f"Error in model-based image analysis: {str(e)}")
                return {
                    'is_property_related': False,
                    'property_confidence': 0.0,
                    'predicted_label': 'error',
                    'top_predictions': [],
                    'image_quality': assess_image_quality(image),
                    'is_ai_generated': False,
                    'authenticity_score': 0.0,
                    'error': str(e)
                }
        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
        }