def __init__(self,
crop=(224, 224),
color=(0.4, 0.4, 0.4, 0.2),
min_area=0.08,
augment=True,
normalize=True,
totensor=True,
num_frames=8,
pad_missing=False,
):
self.crop = crop
self.augment = augment
self.num_frames = num_frames
self.pad_missing = pad_missing
if normalize:
assert totensor
if augment:
transforms = [
video_transforms.RandomResizedCrop(crop, scale=(min_area, 1.)),
video_transforms.RandomHorizontalFlip(),
video_transforms.ColorJitter(*color),
]
else:
transforms = [
video_transforms.Resize(int(crop[0]/0.875)),
video_transforms.CenterCrop(crop),
]
if totensor:
transforms += [volume_transforms.ClipToTensor()]
if normalize:
transforms += [tensor_transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]
self.transform = video_transforms.Compose(transforms)
def __init__(
self,
color=(0.4, 0.4, 0.4, 0.2),
augment=True,
num_frames=8,
pad_missing=False,
random_flip=True,
random_color=True,
):
from utils.videotransforms import video_transforms, volume_transforms, tensor_transforms
self.augment = augment
self.num_frames = num_frames
self.pad_missing = pad_missing
transforms = []
if augment:
if random_flip:
transforms += [video_transforms.RandomHorizontalFlip()]
if random_color:
transforms += [video_transforms.ColorJitter(*color)]
if pad_missing:
transforms += [VideoPadMissing(num_frames)]
transforms += [volume_transforms.ClipToTensor()]
transforms += [
tensor_transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
]
self.transform = video_transforms.Compose(transforms)
def main():
root_dir = "E:\\Datasets\\UCFCrime_img_best"
learning_rate = 1e-3
weight_decay = 0
batch_size = 16
start_epoch = 0
max_epoch = 1
base_size = 8
checkpoint_prefix = 'conv3dae'
train_tfs = video_transforms.Compose([
video_transforms.Resize(224),
video_transforms.RandomCrop(224),
volume_transforms.ClipToTensor()
])
test_tfs = video_transforms.Compose([
video_transforms.Resize(224),
video_transforms.CenterCrop(224),
volume_transforms.ClipToTensor()
])
trainset = UCFCrime(root_dir, train=True, transforms=train_tfs)
testset = UCFCrime(root_dir, train=False, transforms=test_tfs)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=8)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=8)
model = conv3d.Conv3DAE(input_channels=3, base_size=base_size).to(device)
criterion = torch.nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=weight_decay)
checkpoint_path = os.path.join('checkpoints', '{}.tar'.format(checkpoint_prefix))
if os.path.exists(checkpoint_path):
checkpoint = torch.load(checkpoint_path, map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print('Checkpoint loaded, last epoch = {}'.format(checkpoint['epoch'] + 1))
start_epoch = checkpoint['epoch'] + 1
for epoch in range(start_epoch, max_epoch):
train(trainloader, model, optimizer, criterion, None, epoch)
save_model(model, optimizer, epoch, '{}_epoch{}'.format(checkpoint_prefix, epoch + 1))
visualize(testloader, model)
def main():
root_dir = "E:\\Datasets\\UCFCrime_img_best"
output_dir = "E:\\Datasets\\UCFCrime_AE"
batch_size = 16
base_size = 8
checkpoint_prefix = 'conv3dae'
test_tfs = video_transforms.Compose([
video_transforms.Resize(224),
video_transforms.CenterCrop(224),
volume_transforms.ClipToTensor()
])
model = conv3d.Conv3DAE(input_channels=3,
base_size=base_size,
extract_feature=True).to(device)
checkpoint_path = os.path.join('checkpoints',
'{}.tar'.format(checkpoint_prefix))
if os.path.exists(checkpoint_path):
checkpoint = torch.load(checkpoint_path,
map_location=torch.device('cpu'))
model.load_state_dict(checkpoint['state_dict'])
print('Checkpoint loaded, last epoch = {}'.format(checkpoint['epoch'] +
1))
# Train
folder_list = generate_lists(root_dir, train=True)
for row in folder_list:
img_dir, label, img_count = row
extract_feature(model,
root_dir,
os.path.join(output_dir, 'Train'),
img_dir,
label,
img_count,
transforms=test_tfs)
# Test
folder_list = generate_lists(root_dir, train=False)
for row in folder_list:
img_dir, label, img_count = row
extract_feature(model,
root_dir,
os.path.join(output_dir, 'Test'),
img_dir,
label,
img_count,
transforms=test_tfs)