DroneImageClassifier.py

import kagglehub
# Download latest version
path = kagglehub.dataset_download("banuprasadb/visdrone-dataset")

print("Path to dataset files:", path)

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import torchvision
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
import timm
import matplotlib.pyplot as plt # For data viz
import pandas as pd
import numpy as np
import sys
from tqdm.notebook import tqdm
import kagglehub
import random
import time
from ultralytics import YOLO

print('System Version:', sys.version)
print('PyTorch version', torch.__version__)
print('Torchvision version', torchvision.__version__)
print('Numpy version', np.__version__)
print('Pandas version', pd.__version__)

#Class and instantiation
class DroneVisDataset(Dataset):
    def __init__(self, data_dir, transform=None):
        self.data = ImageFolder(data_dir, transform=transform)
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]
    
    @property
    def classes(self):
        return self.data.classes
  
dataset = DroneVisDataset(
    data_dir=path#'/kaggle/input/cards-image-datasetclassification/train'
)

data_dir = path#'/kaggle/input/cards-image-datasetclassification/train'
target_to_class = {
  0: 'pedestrian',
  1: 'people',
  2: 'bicycle',
  3: 'car',
  4: 'van',
  5: 'truck',
  6: 'tricycle',
  7: 'awning-tricycle', #
  8: 'bus',
  9: 'motor'
}
#{v: k for k, v in ImageFolder(data_dir).class_to_idx.items()}
print(target_to_class)

print(f'Total size of the dataset: {len(dataset)}')

rand_index = random.randint(0, len(dataset) - 1)

image, label = dataset[rand_index]

#if you want to grab the image itself grab it exactly from the dataset
                 
print(f'image number {rand_index}')
image

#Let's transform this to a tensor
transform = transforms.Compose([
    transforms.Resize((224, 224)), # originally, this was in shape for 1080 x 1920 pixels, but that had exceptional memory usage
    transforms.ToTensor(),
])

dataset = DroneVisDataset(data_dir, transform)

for image, label in dataset:
    print(image.shape)
    break

dataloader = DataLoader(dataset, batch_size=32, shuffle=True)
for images, labels in dataloader:
  print(images.shape)
  print(labels.shape)
  break

class DroneImageClassifer(nn.Module):
    def __init__(self, num_classes=10):
        super(DroneImageClassifer, self).__init__()


        # This utilizies a torch image model known as efficentnet_b0, which is pretrained
        self.base_model = timm.create_model('efficientnet_b0', pretrained=True) #TODO is this the right model for my image #NOTE consider resnet50
        self.features = nn.Sequential(*list(self.base_model.children())[:-1])

        #and this defines the efficient net size 
        enet_out_size = 1280

        # Make a classifier
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(enet_out_size, num_classes)
        )
    
    def forward(self, x):
        # Connect these parts and return the output
        x = self.features(x)
        output = self.classifier(x)
        return output

model = DroneImageClassifer(num_classes=10)
#print(str(model)[:500])


###USE OF OUR ORIGINAL MODEL, SUBSITUTED FOR YOLO
'''
print("here")
example_out = model(images)
print("but not here")
print(example_out.shape) #batch_size, num_classes

# Loss function
criterion = nn.CrossEntropyLoss()
# Optimizer
optimizer = optim.Adam(model.parameters(), lr=0.001)

criterion(example_out, labels)
print(example_out.shape, labels.shape)
'''
###Training and Validation

'''
train_dir = path + "\VisDrone2019-DET-train"
valid_dir = path + "\VisDrone2019-DET-val"
test_dir = path + "\VisDrone2019-DET-test"
'''

train_img_dir = path + "/VisDrone_Dataset/VisDrone2019-DET-train/images/"
valid_img_dir = path + "/VisDrone_Dataset/VisDrone2019-DET-val/images/"
test_img_dir = path + "/VisDrone_Dataset/VisDrone2019-DET-test-dev/images/"

model = YOLO('yolov8n.pt')

train_results = model(source=train_img_dir, save=True, conf=0.25)
print("successfully loaded training results")
print("moving on to validation")
time.sleep(4)
valid_results = model(source=valid_img_dir, save=True, conf=0.25)
print("successfully loaded validation results")
print("moving on to testing")
time.sleep(4)
test_results = model(source=test_img_dir, save=True, conf=0.25)
print("successfully loaded testing results")


print('here')
train_dataset = DroneVisDataset(train_img_dir, transform=transform)
val_dataset = DroneVisDataset(valid_img_dir, transform=transform)
test_dataset = DroneVisDataset(test_img_dir, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)


#Finally we need to box plot around the images


'''
###

# Simple training loop
num_epochs = 5
train_losses, val_losses = [], []

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model = DroneImageClassifer(num_classes=10)
model.to(device)

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

for epoch in range(num_epochs):
    # Training phase
    model.train()
    running_loss = 0.0
    for images, labels in tqdm(train_loader, desc='Training loop'):
        # Move inputs and labels to the device
        images, labels = images.to(device), labels.to(device)
        
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * labels.size(0)
    train_loss = running_loss / len(train_loader.dataset)
    train_losses.append(train_loss)
    
    # Validation phase
    model.eval()
    running_loss = 0.0
    with torch.no_grad():
        for images, labels in tqdm(val_loader, desc='Validation loop'):
            # Move inputs and labels to the device
            images, labels = images.to(device), labels.to(device)
         
            outputs = model(images)
            loss = criterion(outputs, labels)
            running_loss += loss.item() * labels.size(0)
    val_loss = running_loss / len(val_loader.dataset)
    val_losses.append(val_loss)
    print(f"Epoch {epoch+1}/{num_epochs} - Train loss: {train_loss}, Validation loss: {val_loss}")
    '''