torchvision

import torch from torchvision.transforms import v2 from torchvision import models import os # 1. Create a dummy image tensor (simulating a loaded image) H, W = 256, 256 # Example image dimensions dummy_image = torch.randint(0, 256, size=(3, H, W), dtype=torch.uint8) # 2. Define image transforms using the recommended v2 API preprocess = v2.Compose([ v2.Resize((224, 224), antialias=True), # Resize for common model input sizes v2.ToDtype(torch.float32, scale=True), # Convert to float and scale pixel values to [0, 1] v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), # ImageNet normalization ]) # Apply transforms to the dummy image input_tensor = preprocess(dummy_image) input_batch = input_tensor.unsqueeze(0) # Add a batch dimension (models expect batches) # 3. Load a pre-trained image classification model (e.g., ResNet-18) # Use DEFAULT_WEIGHTS to automatically get the best available pre-trained weights weights = models.ResNet18_Weights.DEFAULT model = models.resnet18(weights=weights) model.eval() # Set the model to evaluation mode for inference # Get the categories the model was trained on for human-readable output categories = weights.meta["categories"] # 4. Perform inference with torch.no_grad(): # Disable gradient calculation for inference to save memory and computations output = model(input_batch) # 5. Get the predicted class probabilities = torch.nn.functional.softmax(output, dim=1) predicted_probability, predicted_idx = torch.max(probabilities, 1) predicted_label = categories[predicted_idx.item()] print(f"Predicted class: {predicted_label} (Probability: {predicted_probability.item():.2f})") print("Quickstart successful: Image processed and classified using torchvision.")