import torch import torch.nn as nn import torch.optim as optim from torchvision import transforms from torch.utils.data import DataLoader, Dataset from fastapi import APIRouter from datetime import datetime from datasets import load_dataset import numpy as np from sklearn.metrics import accuracy_score, precision_score, recall_score import random import os from .utils.evaluation import ImageEvaluationRequest from .utils.emissions import tracker, clean_emissions_data, get_space_info from dotenv import load_dotenv load_dotenv() router = APIRouter() DESCRIPTION = "Random Baseline" ROUTE = "/image" def parse_boxes(annotation_string): """Parse multiple boxes from a single annotation string. Each box has 5 values: class_id, x_center, y_center, width, height""" values = [float(x) for x in annotation_string.strip().split()] boxes = [] # Each box has 5 values for i in range(0, len(values), 5): if i + 5 <= len(values): # Skip class_id (first value) and take the next 4 values box = values[i+1:i+5] boxes.append(box) return boxes def compute_iou(box1, box2): """Compute Intersection over Union (IoU) between two YOLO format boxes.""" # Convert YOLO format (x_center, y_center, width, height) to corners def yolo_to_corners(box): x_center, y_center, width, height = box x1 = x_center - width/2 y1 = y_center - height/2 x2 = x_center + width/2 y2 = y_center + height/2 return np.array([x1, y1, x2, y2]) box1_corners = yolo_to_corners(box1) box2_corners = yolo_to_corners(box2) # Calculate intersection x1 = max(box1_corners[0], box2_corners[0]) y1 = max(box1_corners[1], box2_corners[1]) x2 = min(box1_corners[2], box2_corners[2]) y2 = min(box1_corners[3], box2_corners[3]) intersection = max(0, x2 - x1) * max(0, y2 - y1) # Calculate union box1_area = (box1_corners[2] - box1_corners[0]) * (box1_corners[3] - box1_corners[1]) box2_area = (box2_corners[2] - box2_corners[0]) * (box2_corners[3] - box2_corners[1]) union = box1_area + box2_area - intersection return intersection / (union + 1e-6) def compute_max_iou(true_boxes, pred_box): """Compute maximum IoU between a predicted box and all true boxes""" max_iou = 0 for true_box in true_boxes: iou = compute_iou(true_box, pred_box) max_iou = max(max_iou, iou) return max_iou @router.post(ROUTE, tags=["Image Task"], description=DESCRIPTION) async def evaluate_image(request: ImageEvaluationRequest): """ Evaluate image classification and object detection for forest fire smoke. Current Model: Random Baseline - Makes random predictions for both classification and bounding boxes - Used as a baseline for comparison Metrics: - Classification accuracy: Whether an image contains smoke or not - Object Detection accuracy: IoU (Intersection over Union) for smoke bounding boxes """ # Get space info username, space_url = get_space_info() # Load and prepare the dataset dataset = load_dataset(request.dataset_name, token=os.getenv("HF_TOKEN")) # Split dataset train_test = dataset["train"] test_dataset = dataset["val"] # Start tracking emissions tracker.start() tracker.start_task("inference") #-------------------------------------------------------------------------------------------- # YOUR MODEL INFERENCE CODE HERE # Update the code below to replace the random baseline with your model inference #-------------------------------------------------------------------------------------------- class ImageClassifier(nn.Module): def __init__(self): super(ImageClassifier, self).__init__() self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) self.relu1 = nn.ReLU() self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) self.relu2 = nn.ReLU() self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) self.fc1 = nn.Linear(64 * 16 * 16, 128) self.relu3 = nn.ReLU() self.fc2 = nn.Linear(128, 2) # Output layer with 2 classes (0, 1) def forward(self, x): x = self.pool1(self.relu1(self.conv1(x))) x = self.pool2(self.relu2(self.conv2(x))) x = x.view(x.size(0), -1) x = self.relu3(self.fc1(x)) x = self.fc2(x) return x class CustomDataset(Dataset, labels): def __init__(self, dataset, transform=None): self.dataset = dataset self.transform = transform self.labels = labels def __len__(self): return len(self.dataset) def __getitem__(self, idx): image = self.dataset[idx]['image'] label = self.labels[idx] if self.transform: image = self.transform(image) return image, label # Create an instance of the model model = ImageClassifier() # Define loss function and optimizer criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=0.1) predictions = [] true_labels = [] pred_boxes = [] true_boxes_list = [] # List of lists, each inner list contains boxes for one image # Data Augmentation: torch.manual_seed(0) transform = transforms.Compose([ transforms.RandomCrop(size=(512, 512)), # Crop an image to reduce informations transforms.Resize(size=(64, 64)), # Resize to a standard size, experiment with different sizes transforms.RandomHorizontalFlip(), transforms.RandomVerticalFlip(), transforms.RandomRotation(30), # Add random rotations transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2), # Color variations transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Normalize for ImageNet ]) # Dataset Loader for CNN computation train_loader = DataLoader(train_test, batch_size=64, shuffle=False) test_loader = DataLoader(test_dataset, batch_size=64, shuffle=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) # Training loop num_epochs = 10 for epoch in range(num_epochs): for images, labels in train_loader : images, labels = images.to(device), labels.to(device) # Zero the parameter gradients optimizer.zero_grad() # Forward + backward + optimize outputs = model(images) loss = criterion(outputs, labels) loss.backward() optimizer.step() print(f'Epoch [{epoch + 1}/10], Loss: {loss.item():.4f}') # Evaluation loop model.eval() # Set the model to evaluation mode with torch.no_grad(): for images, labels in test_loader: images, labels = images.to(device), labels.to(device) outputs = model(images) # Apply sigmoid to get probabilities probabilities = torch.sigmoid(outputs) #Get the predicted class with maximum probability _, prediction = torch.max(probabilities, 1) predictions.extend(prediction.cpu().numpy()) for example in test_dataset: # Parse true annotation (YOLO format: class_id x_center y_center width height) annotation = example.get("annotations", "").strip() has_smoke = len(annotation) > 0 true_labels.append(int(has_smoke)) # If there's a true box, parse it and make random box prediction if has_smoke: # Parse all true boxes from the annotation image_true_boxes = parse_boxes(annotation) true_boxes_list.append(image_true_boxes) # For baseline, make one random box prediction per image # In a real model, you might want to predict multiple boxes random_box = [ random.random(), # x_center random.random(), # y_center random.random() * 0.5, # width (max 0.5) random.random() * 0.5 # height (max 0.5) ] pred_boxes.append(random_box) #-------------------------------------------------------------------------------------------- # YOUR MODEL INFERENCE STOPS HERE #-------------------------------------------------------------------------------------------- # Stop tracking emissions emissions_data = tracker.stop_task() # Calculate classification metrics classification_accuracy = accuracy_score(true_labels, predictions) classification_precision = precision_score(true_labels, predictions) classification_recall = recall_score(true_labels, predictions) # Calculate mean IoU for object detection (only for images with smoke) # For each image, we compute the max IoU between the predicted box and all true boxes ious = [] for true_boxes, pred_box in zip(true_boxes_list, pred_boxes): max_iou = compute_max_iou(true_boxes, pred_box) ious.append(max_iou) mean_iou = float(np.mean(ious)) if ious else 0.0 # Prepare results dictionary results = { "username": username, "space_url": space_url, "submission_timestamp": datetime.now().isoformat(), "model_description": DESCRIPTION, "classification_accuracy": float(classification_accuracy), "classification_precision": float(classification_precision), "classification_recall": float(classification_recall), "mean_iou": mean_iou, "energy_consumed_wh": emissions_data.energy_consumed * 1000, "emissions_gco2eq": emissions_data.emissions * 1000, "emissions_data": clean_emissions_data(emissions_data), "api_route": ROUTE, "dataset_config": { "dataset_name": request.dataset_name, "test_size": request.test_size, "test_seed": request.test_seed } } return results