File size: 5,097 Bytes
5cbcc4c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 |
"""
Main file for training Yolo model on Pascal VOC and COCO dataset
"""
import config
import torch
import torch.optim as optim
import albumentations as A
import cv2
from albumentations.pytorch import ToTensorV2
from model import YOLOv3,YOLOV3LITE
from tqdm import tqdm
from utils import (
mean_average_precision,
cells_to_bboxes,
get_evaluation_bboxes,
save_checkpoint,
load_checkpoint,
check_class_accuracy,
get_loaders,
plot_couple_examples
)
from dataset import YOLODatasetOK
from utils import non_max_suppression,plot_image
from loss import YoloLoss
import warnings
warnings.filterwarnings("ignore")
from pytorch_lightning.callbacks import LearningRateMonitor
from pytorch_lightning.callbacks.progress import TQDMProgressBar
from pytorch_lightning.loggers import CSVLogger,TensorBoardLogger
from pytorch_lightning.callbacks import ModelCheckpoint
import torch.optim as optim
import pytorch_lightning as pl
torch.backends.cudnn.benchmark = True
def load_checkpoint(checkpoint_file, model, optimizer, lr):
print("=> Loading checkpoint")
checkpoint = torch.load(checkpoint_file, map_location=config.DEVICE)
model.load_state_dict(checkpoint["state_dict"])
# optimizer.load_state_dict(checkpoint["optimizer"])
# # If we don't do this then it will just have learning rate of old checkpoint
# # and it will lead to many hours of debugging \:
# for param_group in optimizer.param_groups:
# param_group["lr"] = lr
return model
def main():
train_data="100examples.csv"
valid_data = "2examples.csv"
test_data="8examples.csv"
train_loader, test_ldr, train_eval_loader = get_loaders(
train_csv_path=config.DATASET + "/"+train_data, test_csv_path=config.DATASET + "/"+test_data,valid_csv_path=config.DATASET + "/"+valid_data)
#trainer = pl.Trainer()
#test_ldr_final = get_loaders_new(test_csv_path=config.DATASET + "/"+test_data)
trainer= pl.Trainer(
max_epochs=8,
accelerator="auto",
check_val_every_n_epoch=5,
devices=1 if torch.cuda.is_available() else None, # limiting got iPython runs
logger=TensorBoardLogger(save_dir="logs/"),
precision = 16,
#callbacks=[LearningRateMonitor(logging_interval='epoch')]
)
#model = YOLOV3LITE(train_loader,test_ldr,train_eval_loader)
#trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=train_eval_loader)
model_handler = YOLOV3LITE()
# loaded_model =load_checkpoint(
# config.CHECKPOINT_FILE,model_handler.model, model_handler.optimizer, config.LEARNING_RATE
# )
loaded_model =load_checkpoint(
config.CHECKPOINT_FILE,model_handler, model_handler.optimizer, config.LEARNING_RATE
)
test_transform = A.Compose(
[
# Rescale an image so that maximum side is equal to image_size
A.LongestMaxSize(max_size=config.IMAGE_SIZE),
# Pad remaining areas with zeros
A.PadIfNeeded(
min_height=config.IMAGE_SIZE, min_width=config.IMAGE_SIZE, border_mode=cv2.BORDER_CONSTANT
),
# Normalize the image
A.Normalize(
mean=[0, 0, 0], std=[1, 1, 1], max_pixel_value=255
),
# Convert the image to PyTorch tensor
ToTensorV2()
],
# Augmentation for bounding boxes
bbox_params=A.BboxParams(
format="yolo",
min_visibility=0.4,
label_fields=[]
))
dataset = YOLODatasetOK(
csv_file=config.DATASET + "/"+test_data,
img_dir=config.IMG_DIR,
label_dir=config.LABEL_DIR,
S=[13, 26, 52],
anchors=config.ANCHORS,
transform=test_transform
)
# Creating a dataloader object
loader = torch.utils.data.DataLoader(
dataset=dataset,
batch_size=8,
shuffle=True,
)
scaled_anchors = (
torch.tensor(config.ANCHORS)
* torch.tensor(config.S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
).to(config.DEVICE)
plot_couple_examples(loaded_model, loader, 0.6,0.5,scaled_anchors)
# # Defining the grid size and the scaled anchors
# GRID_SIZE = [13, 26, 52]
# scaled_anchors = torch.tensor(config.ANCHORS) / (
# 1 / torch.tensor(GRID_SIZE).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
# )
# # Getting a batch from the dataloader
# x, y = next(iter(loader))
# # Getting the boxes coordinates from the labels
# # and converting them into bounding boxes without scaling
# boxes = []
# for i in range(y[0].shape[1]):
# anchor = scaled_anchors[i]
# boxes += cells_to_bboxes(
# y[i], is_preds=False, S=y[i].shape[2], anchors=anchor
# )[0]
# # Applying non-maximum suppression
# boxes = non_max_suppression(boxes, iou_threshold=1, threshold=0.7)
# # Plotting the image with the bounding boxes
# plot_image(x[0].permute(1,2,0).to("cpu"), boxes)
if __name__ == "__main__":
main() |