Book a Demo

Python SDK & PyTorch Integration

Integrate attention prediction directly into your Python applications with PyTorch

✓ PyTorch Compatible ✓ GPU Accelerated ✓ Easy Integration

Getting Started

Load the model and start predicting attention patterns in your Python environment

📚 Available Models

We offer specialized models for different use cases including Pack AI, Shelf AI, and Visual AI, each available in Basic and Pro versions. Learn more about available models →

Model Loading

import torch
import torchvision.transforms as T
from PIL import Image

# Load the attention prediction model
model = torch.load('attention_model.pt', map_location='cpu')
model.eval()

# Set up image preprocessing
transform = T.Compose([
    T.Resize((512, 512)),
    T.ToTensor(),
    T.Normalize(mean=[0.485, 0.456, 0.406], 
                std=[0.229, 0.224, 0.225])
])

Making Predictions

def predict_attention(image_path):
    # Load and preprocess image
    image = Image.open(image_path).convert('RGB')
    input_tensor = transform(image).unsqueeze(0)
    
    # Generate predictions
    with torch.no_grad():
        predictions = model(input_tensor)
    
    # Extract predicted attention matrices
    hold_matrix = predictions['hold'].squeeze().numpy()
    speed_matrix = predictions['speed'].squeeze().numpy()
    reach_matrix = predictions['reach'].squeeze().numpy()
    
    # Calculate VAI as composite of the three predicted components
    vai_matrix = (hold_matrix + speed_matrix + reach_matrix) / 3.0
    
    return {
        'hold': hold_matrix,
        'speed': speed_matrix,
        'reach': reach_matrix,
        'vai': vai_matrix  # Calculated composite
    }

# Example usage
results = predict_attention('your_image.jpg')

Integration Examples

Batch Processing

import os
from pathlib import Path

def process_image_folder(folder_path):
    results = {}
    for image_file in Path(folder_path).glob('*.{jpg,png,jpeg}'):
        print(f"Processing {image_file.name}...")
        results[image_file.name] = predict_attention(str(image_file))
    
    return results

# Process all images in a directory
batch_results = process_image_folder('./images/')

Heatmap Visualization

import matplotlib.pyplot as plt
import numpy as np

def create_heatmap(attention_matrix, original_image_path):
    # Load original image
    original = Image.open(original_image_path)
    
    # Create heatmap overlay
    plt.figure(figsize=(12, 6))
    
    plt.subplot(1, 2, 1)
    plt.imshow(original)
    plt.title('Original Image')
    plt.axis('off')
    
    plt.subplot(1, 2, 2)
    plt.imshow(original, alpha=0.7)
    plt.imshow(attention_matrix, alpha=0.5, cmap='hot')
    plt.title('Attention Heatmap')
    plt.axis('off')
    
    plt.tight_layout()
    plt.show()

# Visualize results
create_heatmap(results['vai'], 'your_image.jpg')

AOI Analysis

def analyze_aoi(attention_matrix, x, y, width, height):
    """Analyze attention within a specific Area of Interest"""
    
    # Extract AOI region
    aoi_region = attention_matrix[y:y+height, x:x+width]
    
    # Calculate statistics
    mean_attention = np.mean(aoi_region)
    max_attention = np.max(aoi_region)
    total_attention = np.sum(aoi_region)
    
    return {
        'mean_attention': mean_attention,
        'max_attention': max_attention,
        'total_attention': total_attention,
        'coverage': np.count_nonzero(aoi_region) / aoi_region.size
    }

# Analyze specific regions
logo_stats = analyze_aoi(results['vai'], 100, 50, 200, 100)
print(f"Logo attention: {logo_stats}")

Performance Optimization

GPU Acceleration

# Check for GPU availability
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

def predict_attention_gpu(image_path):
    image = Image.open(image_path).convert('RGB')
    input_tensor = transform(image).unsqueeze(0).to(device)
    
    with torch.no_grad():
        predictions = model(input_tensor)
    
    # Move results back to CPU
    return {k: v.cpu().numpy() for k, v in predictions.items()}

Memory Management

# Clear GPU memory between predictions
import gc

def predict_with_cleanup(image_path):
    result = predict_attention_gpu(image_path)
    
    # Clear GPU cache
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    
    # Force garbage collection
    gc.collect()
    
    return result

Ready to Start Developing?

Get access to our PyTorch models and start building attention-aware applications

Request Model Access View API Alternative