50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
import os
|
|
from torch import tensor
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
from torchvision.transforms import ToTensor
|
|
import torch.nn as nn
|
|
|
|
from .r_cnn import ObjectDetectionCNN, train, eval
|
|
from .cocoDetectionDataset import CocoDetectionDataset
|
|
|
|
|
|
def get_transform():
|
|
return ToTensor()
|
|
|
|
def train_cnn_test():
|
|
|
|
SAVE_PATH = "./saved_models"
|
|
PADDING = 20
|
|
|
|
train_dataset = CocoDetectionDataset(
|
|
image_dir="data/football/train",
|
|
annotation_path="data/football/train/_annotations.coco.json",
|
|
transforms=get_transform()
|
|
)
|
|
|
|
val_dataset = CocoDetectionDataset(
|
|
image_dir="data/football/valid",
|
|
annotation_path="data/football/valid/_annotations.coco.json",
|
|
transforms=get_transform()
|
|
)
|
|
|
|
train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
|
|
val_loader = DataLoader(val_dataset, batch_size=2, shuffle=True, collate_fn=lambda x: tuple(zip(*x)))
|
|
|
|
device = torch.device("cpu") if not torch.cuda.is_available() else torch.device("cuda:0")
|
|
print("Using device", device)
|
|
|
|
model = ObjectDetectionCNN(c_in=3, c_hidden=32, c_out=2, layers=10)
|
|
model.to(device)
|
|
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-4)
|
|
loss_module = nn.CrossEntropyLoss()
|
|
|
|
#train(model=model, loss_module=loss_module, train_loader=train_loader, val_loader=val_loader,
|
|
# optimizer=optimizer, SAVE_PATH=SAVE_PATH, saving=True, PADDING=40, device=device)
|
|
|
|
#exit()
|
|
|
|
image = next(iter(val_loader))[0][0]
|
|
eval(model=model, image=image, BUILD_PATH=os.path.join(SAVE_PATH, "object_detection", "object_detection"),
|
|
device=device, PADDING=40, minSize=5, maxSize=100, minConf=0.8) |