class ResNet50_FasterRCNN:
def __init__(self, pretrained=False):
# Building our FasterRCNN model for objects detection
backbone = resnet_fpn_backbone('resnet50', pretrained=pretrained)
num_classes = 4 + 1
anchor_generator = AnchorGenerator(sizes=(40, 60, 150, 200, 250),
aspect_ratios=(0.7, 1.0, 1.3))
self.model = FRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator)
def train(self):
self.model.train()
def to(self, device):
self.model.to(device)
def eval(self):
self.model.eval()
def parameters(self):
return self.model.parameters()
def get_state_dict(self):
return self.model.state_dict()
def set_state_dict(self, state_dict):
self.model.load_state_dict(state_dict)
def fit_batch(self, images, target):
return self.model(images, target)
def predict_batch(self, images):
return self.model(images)
class TorchDetector:
"""
Torch object detector
"""
def __init__(self, config, logger):
self._logger = logger
self._threshold = config['threshold']
modelfile = config['model']
self._device = config['device'] # cpu, cuda, cuda:0
backbone = resnet_fpn_backbone('resnet50', False)
self._model = FasterRCNN(backbone, 8) # 8 classes
checkpoint = torch.load(modelfile, map_location=self._device)
self._model.load_state_dict(checkpoint['model_state_dict'])
device = torch.device(self._device)
self._model.to(device)
self._model.eval()
def stop(self):
"""
Destruction
"""
def detectObjects(self, img) -> List[e.DetectedObject]:
"""
Implementation of detector interface
"""
wsize = 1600
hsize = 800
_pretransform = A.Compose([
A.Resize(hsize, wsize),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2(),
])
image_tensor = _pretransform(image=img)['image']
tstart = time.time()
outputs = self._model.forward(
image_tensor.unsqueeze(0).float().to(device=self._device))
classes = outputs[0]['labels'].detach().cpu().numpy()
scores = outputs[0]['scores'].detach().cpu().numpy()
boxes = outputs[0]['boxes'].detach().cpu().numpy()
self._logger.debug(
f'Torch model inferring time: {time.time() - tstart}')
result = zip(classes, scores, boxes)
h, w, _ = img.shape
wscale = w / wsize
hscale = h / hsize
#print(f'h,w:{h},{w}; wsc,hsc:{wscale},{hscale}')
#print(list(result))
return ObjectDetector.getDetectedObjectsCollection(
result, hscale, wscale, self._threshold, False)
class Detect:
def __init__(self):
super().__init__()
backbone = torchvision.models.vgg16(pretrained=False).features
backbone.out_channels = 512
anchor_sizes = ((8, 16, 32, 64, 128, 256, 512), )
aspect_ratios = ((1 / 2, 1 / 3, 1 / 4, 1 / 5, 1 / 6, 1 / math.sqrt(2),
1, 2, math.sqrt(2), 3, 4, 5, 6, 7, 8), )
anchor_generator = AnchorGenerator(sizes=anchor_sizes,
aspect_ratios=aspect_ratios)
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=['0', '1', '2', '3', '4'],
output_size=7,
sampling_ratio=2)
self.model = FasterRCNN(backbone,
num_classes=7,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
self.device = torch.device('cpu')
self.model.load_state_dict(torch.load('2.pth'))
self.model.to(self.device)
self.model.eval()
def forward(self, img):
img = torch.tensor(img, dtype=torch.float32) / 255
img = img.permute((2, 0, 1))
output = model([img.to(self.device)])
boxes = output[0]['boxes']
labels = output[0]['labels']
scores = output[0]['scores']
last = {}
result = {}
for i, v in enumerate(labels):
if v == 1 and scores[i] > last['send']:
last['send'] = scores[i]
result['send'] = boxes[i]
elif v == 2 and scores[i] > last['number']:
last['number'] = scores[i]
result['number'] = boxes[i]
elif v == 3 and scores[i] > last['date']:
last['date'] = scores[i]
result['date'] = boxes[i]
elif v == 4 and scores[i] > last['quote']:
last['quote'] = scores[i]
elif v == 5 and scores[i] > last['header']:
last['header'] = scores[i]
result['header'] = boxes[i]
elif v == 6 and scores[i] > last['motto']:
last['motto'] = scores[i]
result['motto'] = boxes[i]
# elif v == 7 and scores[i] > last['secrete']:
# last['secrete'] = scores[i]
# result['secrete'] = boxes[i]
# elif v == 8 and scores[i] > last['sign']:
# last['sign'] = scores[i]
# result['sign'] = boxes[i]
return result
def train_mask():
torch.manual_seed(1)
dataset = unimib_data.UNIMIBDataset(get_transform(train=True))
dataset_test = unimib_data.UNIMIBDataset(get_transform(train=False))
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=1, shuffle=True,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=1, shuffle=False,
collate_fn=utils.collate_fn)
# device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
device = "cpu"
num_classes = 74
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=True)
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
for param in backbone:
param.requires_grad = False
backbone.out_channels = 1280
model.backbone = backbone
model = FasterRCNN(backbone,
num_classes=74)
model = torch.nn.Sequential(
model.rpn,
model.roi_heads
)
print(model)
# for param in model.backbone:
# param.requires_grad = False
# model = get_instance_segmentation_model(num_classes)
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
momentum=0.9, weight_decay=0.0005)
# and a learning rate scheduler which decreases the learning rate by
# 10x every 3 epochs
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
num_epochs = 1
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
# train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
# update the learning rate
# lr_scheduler.step()
# evaluate on the test dataset
evaluate(model, data_loader_test, device=device)
def model_predictions_from_loader(
model: FasterRCNN,
data_loader: DataLoader,
device: torch.device,
max_n_batches=1,
normalization_transform: Normalize = Normalize([1, 1, 1], [1, 1, 1])
) -> List[Tuple[BoxedExample, BoxedExample]]:
model = model.to(device).eval()
input_examples = []
output_examples = []
with torch.no_grad():
for i, (img, labels) in enumerate(data_loader):
if i >= max_n_batches:
break
outputs = model([i.to(device) for i in img])
input_exs = [
model_input_to_boxed_example((im, l), normalization_transform)
for (im, l) in zip(img, labels)
]
output_exs = [
i.replace(boxes=o["boxes"].cpu().numpy())
for (i, o) in zip(input_exs, outputs)
]
input_examples.extend(input_exs)
output_examples.extend(output_exs)
return list(zip(input_examples, output_examples))
def model_creation(pretrain=True,
num_classes=5,
num_epoch=20,
device="cuda:0"):
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=True)
num_classes = 5
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
if device == 'cuda:0':
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
if device == "cpu":
print("cuda is no available")
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
# let's train it for 10 epochs
return model, optimizer, lr_scheduler, num_epoch,
def evaluate_for_losses(model: FasterRCNN,
data_loader: DataLoader,
device: torch.device,
max_n_batches: int = 1) -> Dict[str, float]:
model = model.to(device)
model.train()
loss_dicts = []
with torch.no_grad():
for i, (img, labels) in enumerate(tqdm(data_loader)):
if i >= max_n_batches:
break
img = list(image.to(device) for image in img)
labels = [{k: v.to(device) for k, v in t.items()} for t in labels]
loss_dicts.append(model(img, labels))
return {
k: np.mean([ld[k].cpu().item() for ld in loss_dicts])
for k in loss_dicts[0].keys()
}
class FasterRCNNFood:
def __init__(self,
backbone_name: str,
pretrained: bool = True,
finetune: bool = True,
num_classes: int = 2):
self.__pretrained = pretrained
self.__num_classes = num_classes
self.__model_name = backbone_name
backbone = build_backbone(backbone_name, pretrained, finetune)
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
self.model = FasterRCNN(backbone=backbone,
num_classes=num_classes,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
self.params = [p for p in self.model.parameters() if p.requires_grad]
self.optimizer = torch.optim.Adam(params=self.params,
lr=0.005,
weight_decay=0.0005)
self.lr_scheduler = torch.optim.lr_scheduler.StepLR(
optimizer=self.optimizer, step_size=3, gamma=0.1)
def train(self,
data_loader: DataLoader,
data_loader_test: DataLoader,
num_epochs: int = 10,
use_cuda: bool = True,
epoch_save_ckpt: Union[int, list] = None,
dir: str = None):
"""
Method to train FasterRCNNFood model.
Args:
data_loader (torch.utils.data.DataLoader): data loader to train model on
data_loader_test (torch.utils.data.DataLoader): data loader to evaluate model on
num_epochs (int = 10): number of epoch to train model
use_cuda (bool = True): use cuda or not
epoch_save_ckpt (list or int): Epoch at which you want to save the model. If -1 save only last epoch.
dir (str = "models/): Directory where model are saved under the name "{model_name}_{date}_ep{epoch}.pth"
"""
if epoch_save_ckpt == -1:
epoch_save_ckpt = [num_epochs - 1]
if not dir:
dir = "models"
dir = Path(dir)
dir.mkdir(parents=True, exist_ok=True)
# choose device
if use_cuda and torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
# define dataset
self.model.to(device)
writer = SummaryWriter()
for epoch in range(num_epochs):
# train for one epoch, printing every 50 iterations
train_one_epoch(self.model,
self.optimizer,
data_loader,
device,
epoch,
print_freq=50,
writer=writer)
# update the learning rate
self.lr_scheduler.step()
# evaluate on the test dataset
evaluate(self.model,
data_loader_test,
device=device,
writer=writer,
epoch=epoch)
# save checkpoint
if epoch in epoch_save_ckpt:
self.save_checkpoint(dir.as_posix(), epoch)
writer.close()
print("That's it!")
def save_checkpoint(self, dir: str, epoch: int):
"""
Save a model checkpoint at a given epoch.
Args:
dir: dir folder to save the .pth file
epoch: epoch the model is
"""
state = {
'epoch': epoch + 1,
'state_dict': self.model.state_dict(),
'optimizer': self.optimizer.state_dict(),
'num_classes': self.__num_classes,
'pretrained': self.__pretrained,
"model_name": self.__model_name
}
now = datetime.now()
filename = "{model_name}_{date}_ep{epoch}.pth".format(
model_name=self.__model_name,
date=now.strftime("%b%d_%H-%M"),
epoch=epoch)
torch.save(state, Path(dir) / filename)
"Checkpoint saved : {}".format(Path(dir) / filename)
def predict(self, dataset, idx):
img, _ = dataset[idx]
img.to("cpu")
self.model.eval()
self.model.to("cpu")
pred = self.model([img])
return img, pred[0]
@staticmethod
def load_checkpoint(filename: str,
cuda: bool = True) -> ("FasterRCNNFood", int):
"""
Load a model checkpoint to continue training.
Args:
filename (str): filename/path of the checkpoint.pth
cuda (bool = True): use cuda
Returns:
(FasterRCNNFood) model
(int) number of epoch + 1 the model was trained with
"""
device = torch.device("cuda") if (
cuda and torch.cuda.is_available()) else torch.device("cpu")
start_epoch = 0
if Path(filename).exists():
print("=> loading checkpoint '{}'".format(filename))
checkpoint = torch.load(filename, map_location=device)
# Load params
pretrained = checkpoint['pretrained']
num_classes = checkpoint["num_classes"]
start_epoch = checkpoint['epoch']
model_name = checkpoint['model_name']
# Build model key/architecture
model = FasterRCNNFood(model_name, pretrained, num_classes)
# Update model and optimizer
model.model.load_state_dict(checkpoint['state_dict'])
model.optimizer.load_state_dict(checkpoint['optimizer'])
model.model = model.model.to(device)
# now individually transfer the optimizer parts...
for state in model.optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.to(device)
print("=> loaded checkpoint '{}' (epoch {})".format(
filename, checkpoint['epoch']))
return model, start_epoch
else:
print("=> no checkpoint found at '{}'".format(filename))
@staticmethod
def load_for_inference(filename: str,
cuda: bool = True) -> "FasterRCNNFood":
"""
Load a model checkpoint to make inference.
Args:
filename (str): filename/path of the checkpoint.pth
cuda (bool = True): use cuda
Returns:
(FasterRCNNFood) model
"""
device = torch.device("cuda") if (
cuda and torch.cuda.is_available()) else torch.device("cpu")
if Path(filename).exists():
print("=> loading checkpoint '{}'".format(filename))
checkpoint = torch.load(filename, map_location=device)
# Load params
pretrained = checkpoint['pretrained']
num_classes = checkpoint["num_classes"]
model_name = checkpoint['model_name']
# Build model key/architecture
model = FasterRCNNFood(model_name, pretrained, num_classes)
# Update model and optimizer
model.model.load_state_dict(checkpoint['state_dict'])
model.model = model.model.to(device)
model.model = model.model.eval()
print("=> loaded checkpoint '{}'".format(filename))
return model
else:
print("=> no checkpoint found at '{}'".format(filename))
shuffle=False,
num_workers=4,
collate_fn=utils.collate_fn)
# In[10]:
device = torch.device('cuda') if torch.cuda.is_available() else torch.device(
'cpu')
# our dataset has two classes only - background and person
num_classes = 2
# get the model using our helper function
model = get_instance_segmentation_model(num_classes)
# move model to the right device
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
# and a learning rate scheduler which decreases the learning rate by
# 10x every 3 epochs
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
# In[11]:
sim_clr.load_state_dict(torch.load(sim_dict, map_location=device))
model_res34 = torchvision.models.resnet34(pretrained=False)
network_helpers.copy_weights_between_models(sim_clr, model_res34)
modules = list(model_res34.children())[:-1]
backbone = nn.Sequential(*modules)
backbone.out_channels = 512
model_fnn = FasterRCNN(backbone=backbone, num_classes=10)
# ImageNet
#in_features = model_fnn.roi_heads.box_predictor.cls_score.in_features
#model_fnn.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes=10)
model_fnn.to(device)
params_fnn = [p for p in model_fnn.parameters() if p.requires_grad]
params_cnn_fnn = list(params_cnn) + list(params_fnn)
optimizer = optim.SGD(params_cnn_fnn,
lr=args.lr,
momentum=0.6,
weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer,
base_lr=1e-3,
max_lr=6e-3)
# lr_scheduler = CosineAnnealingLR(optimizer, 50, eta_min=1e-3, last_epoch=-1)
model_save_path = args.save_final_model_path
num_workers=cfg.workers,
collate_fn=collate_fn)
# Start to traning FASTER RCNN
backbone = backboneNet_efficient() # use efficientnet as our backbone
backboneFPN = backboneWithFPN(backbone) # add FPN
anchor_generator = AnchorGenerator(cfg.anchor_sizes, cfg.aspect_ratios)
model_ft = FasterRCNN(backboneFPN,
num_classes=cfg.num_classes,
rpn_anchor_generator=anchor_generator,
min_size=cfg.min_size,
max_size=cfg.max_size)
model_ft.to(device)
optimizer_ft = optim.SGD(model_ft.parameters(),
lr=cfg.learning_rate,
momentum=cfg.momentum,
weight_decay=cfg.weight_decay)
lr_scheduler = lr_scheduler.MultiStepLR(optimizer_ft,
milestones=cfg.milestones,
gamma=cfg.gamma)
model_ft = train_model(model_ft,
train_loader,
valid_loader,
optimizer_ft,
lr_scheduler,
def main():
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(
pretrained=True)
num_classes = 2
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=2,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# Train classifier
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
dataset = OysterDataset(
"/Users/darwin/Downloads/Bay project/Training",
"/Users/darwin/Downloads/Bay project/Training/TrainingDetectionLabels2.csv",
get_transform(train=True))
dataset_test = OysterDataset(
"/Users/darwin/Downloads/Bay project/Training",
"/Users/darwin/Downloads/Bay project/Training/TrainingDetectionLabels2.csv",
get_transform(train=False))
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:100])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=2,
shuffle=False,
num_workers=4,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(dataset_test,
batch_size=2,
shuffle=False,
num_workers=4,
collate_fn=utils.collate_fn)
model = get_model_instance_segmentation(num_classes)
model.to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.8,
weight_decay=0.0005)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
num_epochs = 10
for epoch in range(num_epochs):
train_one_epoch(model,
optimizer,
data_loader,
device,
epoch,
print_freq=10)
lr_scheduler.step()
evaluate(model, data_loader_test, device=device)
print("That's it!")
class FasterRCNNMODEL:
#TODO: Later on enable passing params params
def __init__(self, model_params=None):
self.params = model_params
self.model = None
self.optimizer = None
self.device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
def set_backbone(self, backbone):
"""
backbone is a string containing the backbone we want to use in the model. add more options
"""
if 'vgg' in backbone.lower():
"to somthing-check for options"
elif 'mobilenet_v2' in backbone.lower():
self.backbone = torchvision.models.mobilenet_v2(
pretrained=True).features
self.backbone.out_channels = 1280
elif 'resnet50' in backbone.lower():
self.backbone = torchvision.models.resnet50(
pretrained=True).features
self.backbone.out_channels = 256
def set_model(self):
"""
Set model and determine configuration
:return: None, generate self.model to be used for training and testing
"""
# Default values: box_score_thresh = 0.05, box_nms_thresh = 0.5
kwargs = {
'box_score_thresh': 0.3,
'box_nms_thresh': 0.3,
'box_detections_per_img': 6
}
# self.model = torchvision.models.detection.fasterrcnn_resnet50_fpn(pretrained=False,
# pretrained_backbone=True,
# **kwargs)
self.model = FasterRCNN(self.backbone, num_classes=7, **kwargs)
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
num_classes = 7
in_features = self.model.roi_heads.box_predictor.cls_score.in_features
self.model.roi_heads.box_predictor = FastRCNNPredictor(
in_features, num_classes)
# Allow Multiple GPUs:
# if torch.cuda.device_count() > 1:
# self.model = nn.DataParallel(self.model)
self.model = self.model.to(device)
if self.params is None:
params = [p for p in self.model.parameters() if p.requires_grad]
else:
# TODO: Enable user defined model params
pass
self.optimizer = torch.optim.SGD(params, lr=0.01)
def train_model(self, train_loader, num_epochs):
"""
Train (only!) of the model
:param train_loader: DataLoader object
:param num_epochs: int. Number of epochs to train the model
:return: None,
"""
self.model.train() # Set to training mode
for epoch in range(num_epochs):
for images, targets in train_loader:
images = list(image.to(self.device) for image in images)
targets = [{k: v.to(self.device)
for k, v in t.items()} for t in targets]
# Zero Gradients
self.optimizer.zero_grad()
# self.model = self.model.double()
# Calculate Loss
loss_dict = self.model(images, targets) # what happens here?
losses = sum(loss for loss in loss_dict.values())
losses.backward()
# Update weights
self.optimizer.step()
print('Train Loss = {:.4f}'.format(losses.item()))
def train_eval_model(self, train_loader, val_loader, num_epochs):
"""
Train model and evaluate performance after each epoch
:param train_loader: DataLoader object. Training images and targets
:param val_loader: DataLoader object. validation images and targets
:param num_epochs: int. Number of epochs for training and validation
:return:
"""
# For evaluation
imgs_name_list = []
bbox_list = []
labels_list = []
for epoch in range(num_epochs):
train_loss = 0
val_loss = 0
self.model.train() # Set to training mode
with torch.set_grad_enabled(True):
for images, targets in train_loader:
# Pass data to GPU
images = list(image.to(self.device) for image in images)
targets = [{k: v.to(self.device)
for k, v in t.items()} for t in targets]
# Zero Gradients
self.optimizer.zero_grad()
# self.model = self.model.double()
# Calculate Loss
loss_dict = self.model(images,
targets) # what happens here?
losses = sum(loss for loss in loss_dict.values())
train_loss += losses.item() * len(images)
# Backward Prop & Update weights
losses.backward()
self.optimizer.step()
print('Train Loss = {:.4f}'.format(train_loss /
len(train_loader.dataset)))
# TODO: Calculate Dice and IoU loss for it
with torch.no_grad():
for idx, (imgs_name, images, targets) in enumerate(val_loader):
self.model.train()
images = list(image.to(self.device) for image in images)
targets = [{k: v.to(self.device)
for k, v in t.items()} for t in targets]
loss_dict = self.model(images, targets)
losses = sum(loss for loss in loss_dict.values())
val_loss += losses.item() * len(images)
if epoch == num_epochs - 1:
self.model.eval() # Set model to evaluate performance
targets = self.model(images)
# Think of moving all this into gen_out_file - Looks nicer
imgs_name_list.extend(imgs_name)
bbox_list.extend([
target['boxes'].int().cpu().tolist()
for target in targets
])
labels_list.extend([
target['labels'].int().cpu().tolist()
for target in targets
])
"""Optional - SEE the performance on the second last batch"""
if (epoch == num_epochs - 1) and idx == (len(val_loader) -
2):
self.model.eval() # Set model to evaluate performance
targets = self.model(images)
MiscUtils.view(images,
targets,
k=len(images),
model_type='faster_rcnn')
DataUtils.gen_out_file('output_file.txt', imgs_name_list,
bbox_list, labels_list)
print('Validation Loss = {:.4f}'.format(
val_loss / len(val_loader.dataset)))
def main():
parser = argparse.ArgumentParser(description='VISUM 2019 competition - baseline training script', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-d', '--data_path', default='/home/master/dataset/train', metavar='', help='data directory path')
parser.add_argument('-m', '--model_path', default='./baseline.pth', metavar='', help='model file (output of training)')
parser.add_argument('--epochs', default=50, type=int, metavar='', help='number of epochs')
parser.add_argument('--lr', default=0.005, type=float, metavar='', help='learning rate')
parser.add_argument('--l2', default=0.0005, type=float, metavar='', help='L-2 regularization')
args = vars(parser.parse_args())
# Data augmentation
def get_transform(train):
transforms = []
# converts the image, a PIL image, into a PyTorch Tensor
transforms.append(T.ToTensor())
if train:
# during training, randomly flip the training images
# and ground-truth for data augmentation
transforms.append(T.RandomHorizontalFlip(0.5))
transforms.append(torchvision.transforms.ColorJitter(contrast=.3))
transforms.append(torchvision.transforms.ColorJitter(brightness=.4))
return T.Compose(transforms)
backbone = torchvision.models.mobilenet_v2(pretrained=True).features
backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512),),
aspect_ratios=((0.5, 1.0, 2.0),))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=[0],
output_size=7,
sampling_ratio=2)
# put the pieces together inside a FasterRCNN model
model = FasterRCNN(backbone,
num_classes=10,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# See the model architecture
print(model)
# use our dataset and defined transformations
dataset = VisumData(args['data_path'], modality='rgb', transforms=get_transform(train=True))
dataset_val = VisumData(args['data_path'], modality='rgb', transforms=get_transform(train=False))
# split the dataset in train and test set
torch.manual_seed(1)
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-100])
dataset_val = torch.utils.data.Subset(dataset_val, indices[-100:])
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=2, shuffle=True, num_workers=0,
collate_fn=utils.collate_fn)
data_loader_val = torch.utils.data.DataLoader(
dataset_val, batch_size=2, shuffle=False, num_workers=0,
collate_fn=utils.collate_fn)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=args['lr'],
momentum=0.9, weight_decay=args['l2'])
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=10,
gamma=0.5)
for epoch in range(args['epochs']):
# train for one epoch, printing every 10 iterations
epoch_loss = train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=nvid)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
evaluator = evaluate(model, data_loader_val, device=device)
torch.save(model, args['model_path'])
def main(args):
detection_util.init_distributed_mode(args)
print(args)
device = torch.device(args.device)
torch.multiprocessing.set_sharing_strategy('file_system')
#train loader
mean = (0.4914, 0.4822, 0.4465)
std = (0.2023, 0.1994, 0.2010)
normalize = transforms.Normalize(mean=mean, std=std)
train_transform = transforms.Compose([transforms.ToTensor(), normalize])
train_dataset = SVHNFull(root='./datasets/SVHN_full',
transform=train_transform,
download=False)
test_dataset = SVHNFull(root='./datasets/SVHN_full',
split='test',
transform=train_transform,
download=False)
if args.distributed:
train_sampler = data.distributed.DistributedSampler(train_dataset)
test_sampler = data.distributed.DistributedSampler(test_dataset)
else:
train_sampler = None
test_sampler = None
print(args.batch_size)
train_loader = data.DataLoader(train_dataset,
batch_size=args.batch_size,
shuffle=(train_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=train_sampler,
collate_fn=train_dataset.collate_fn)
test_loader = data.DataLoader(test_dataset,
batch_size=args.batch_size,
shuffle=(test_sampler is None),
num_workers=args.workers,
pin_memory=True,
sampler=test_sampler,
collate_fn=test_dataset.collate_fn)
pre_mod = ResNetDetection(name='resnet18')
#pre_mod = SupResNetDetection(name='resnet18')
if args.pretrained:
ckpt = torch.load(args.pretrained, map_location='cpu')
state_dict = ckpt['model']
new_state_dict = {}
for k, v in state_dict.items():
k = k.replace("module.", "")
new_state_dict[k] = v
state_dict = new_state_dict
pre_mod.load_state_dict(state_dict)
backbone = pre_mod.encoder
backbone.out_channels = 512
anchor_generator = AnchorGenerator(sizes=((16, 32, 128), ),
aspect_ratios=((0.5, 1.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],
output_size=3,
sampling_ratio=2)
model = FasterRCNN(backbone,
num_classes=10,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler,
rpn_batch_size_per_image=64)
model.to(device)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)
if args.resume:
checkpoint = torch.load(args.resume, map_location='cpu')
model_without_ddp.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
args.start_epoch = checkpoint['epoch'] + 1
print("Start training")
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
train_sampler.set_epoch(epoch)
train(model, optimizer, train_loader, device, epoch, args.print_freq)
lr_scheduler.step()
#if epoch > 15:
evaluate(model, test_loader, device=device)
detection_util.save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'args': args,
'epoch': epoch
}, os.path.join('save/detector/', 'model_{}.pth'.format(epoch)))
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
def main(network):
# train on the GPU or on the CPU, if a GPU is not available
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
# our dataset has two classes only - background and person
num_classes = 2
#dataset = torch.utils.data.Subset(TBdata,[range(len(TBdata))])
indices = torch.randperm(len(TBdata)).tolist()
dataset = torch.utils.data.Subset(TBdata, indices[:])
indices_ = torch.randperm(len(TBdata_test)).tolist()
dataset_val = torch.utils.data.Subset(TBdata_test, indices_[:])
# get the model using our helper function
#model = get_model_instance_segmentation(num_classes)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=8,
sampler=None,
num_workers=0,
collate_fn=collate_fn)
dataloader_val = torch.utils.data.DataLoader(dataset_val,
batch_size=8,
sampler=None,
num_workers=0,
collate_fn=collate_fn)
#Calculated statistics on training data:
#Transform parameters
min_size = 550
max_size = 700
image_means = [0.9492, 0.9492, 0.9492]
image_stds = [0.1158, 0.1158, 0.1158]
if network == 'resnet50':
backbone = resnet_fpn_backbone('resnet50', True)
model = FasterRCNN(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'resnet18':
backbone = resnet_fpn_backbone('resnet18', True)
model = FasterRCNN(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'resnet152':
backbone = resnet_fpn_backbone('resnet152', True)
model = FasterRCNN(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'RPNresnet50':
backbone = resnet_fpn_backbone('resnet50', True)
model = RPN_custom(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
elif network == 'RPNresnet152':
backbone = resnet_fpn_backbone('resnet152', True)
model = RPN_custom(backbone,
num_classes,
min_size=min_size,
max_size=max_size,
image_mean=image_means,
image_std=image_stds)
# move model to the right device
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params,
lr=0.005,
momentum=0.9,
weight_decay=0.0005)
# and a learning rate scheduler
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
num_epochs = 10
Ls = {
'total loss': [],
'loss_classifier': [],
'loss_box_reg': [],
'loss_objectness': [],
'loss_rpn_box_reg': []
}
Ls_val = {
'total loss': [],
'loss_classifier': [],
'loss_box_reg': [],
'loss_objectness': [],
'loss_rpn_box_reg': []
}
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
train_one_epoch(model,
optimizer,
dataloader,
device,
epoch,
print_freq=10)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
#evaluate(model, dataloader_test, device=device)
Ls_val = record_losses(model, dataloader_val, device, Ls_val, network)
#record losses
Ls = record_losses(model, dataloader, device, Ls, network)
#If folder does not exist already, create it
output_loc = f'./{network}/'
if not os.path.exists(output_loc):
os.makedirs(output_loc)
torch.save(model.state_dict(), output_loc + 'model.pt')
print("That's it!")
return Ls, Ls_val, num_epochs
transform_img = transforms.Compose([transforms.ToTensor()])
EPOCH = 250
CLASSES = 3
DEVICE = torch.device("cuda")
BATCH_SIZE = 10
anchor_generator = AnchorGenerator(sizes=((32, 64), ),
aspect_ratios=((0.6, 1.0, 1.6), ))
backbone = torchvision.models.vgg19(pretrained=False).features
backbone.out_channels = 512
model = FasterRCNN(backbone,
num_classes=CLASSES,
rpn_anchor_generator=anchor_generator)
model.load_state_dict(
torch.load('models_new/' + 'model_' + str(EPOCH) + '.pth'))
model.to(DEVICE)
model.eval()
start_time = time.time()
ear_count = 0
for T in types:
for E in ears:
CTs = os.listdir(data_path + dataset_name + T + E)
for CT in CTs:
print('current path:{}'.format(data_path + dataset_name + T + E +
CT))
ear_count += 1
img_names = glob.glob(data_path + dataset_name + T + E + CT +
'/*.jpg')
sorted(img_names, key=lambda x: x.split('\\')[-1])
with torch.no_grad():
start, end = 0, BATCH_SIZE
# model = torchvision.models.mobilenet_v2(pretrained=True)
# backbone = model.features
# backbone.out_channels = 1280
anchor_generator = AnchorGenerator(sizes=((32, 64, 128, 256, 512), ),
aspect_ratios=((0.5, 1.0, 2.0), ))
roi_pooler = torchvision.ops.MultiScaleRoIAlign(
featmap_names=[0, 1, 2, 3, 4, 5], output_size=7, sampling_ratio=2)
model_RCNN = FasterRCNN(backbone,
num_classes=19,
rpn_anchor_generator=anchor_generator,
box_roi_pool=roi_pooler)
# model_RCNN = torch.load("model_RCNN_mobilenet.pth")
model_RCNN = model_RCNN.to(device)
def collate_fn(batch):
return tuple(zip(*batch))
dataset = exp4Dataset("./image_exp/Detection",
T.Compose([T.RandomHorizontalFlip(0.5),
T.ToTensor()]))
data_loader = torch.utils.data.DataLoader(dataset,
batch_size=2,
shuffle=True,
num_workers=4,
collate_fn=collate_fn)