Update tasks/image.py

tasks/image.py

@@ -1,3 +1,9 @@
+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
@@ -100,21 +106,113 @@ async def evaluate_image(request: ImageEvaluationRequest):
     # 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))
-        
-        # Make random classification prediction
-        pred_has_smoke = random.random() > 0.5
-        predictions.append(int(pred_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