Full Code of JakubSochor/BoxCars for AI

Repository: JakubSochor/BoxCars
Branch: master
Commit: 8757b8b36857
Files: 15
Total size: 20.6 MB

Directory structure:
gitextract_5ybgmjqh/

├── .gitignore
├── README.md
├── data/
│   └── estimated_3DBB.pkl
├── lib/
│   ├── __init__.py
│   ├── boxcars_data_generator.py
│   ├── boxcars_dataset.py
│   ├── boxcars_image_transformations.py
│   └── utils.py
├── models/
│   ├── .gitignore
│   └── README.md
├── requirements.txt
└── scripts/
    ├── _init_paths.py
    ├── config.py
    ├── download_models.py
    └── train_eval.py

================================================
FILE CONTENTS
================================================

================================================
FILE: .gitignore
================================================
cache


================================================
FILE: README.md
================================================
# BoxCars Fine-Grained Recognition of Vehicles
This is Keras+Tensorflow re-implementation of our method for fine-grained classification of vehicles decribed in **BoxCars: Improving Vehicle Fine-Grained Recognition using 3D Bounding Boxes in Traffic Surveillance** ([link](https://doi.org/10.1109/TITS.2018.2799228)).
The numerical results are slightly different, but similar. This code is for **research only** purposes.
If you use the code, please cite our paper:
```
@ARTICLE{Sochor2018, 
author={J. Sochor and J. Špaňhel and A. Herout}, 
journal={IEEE Transactions on Intelligent Transportation Systems}, 
title={BoxCars: Improving Fine-Grained Recognition of Vehicles Using 3-D Bounding Boxes in Traffic Surveillance}, 
year={2018}, 
volume={PP}, 
number={99}, 
pages={1-12}, 
doi={10.1109/TITS.2018.2799228}, 
ISSN={1524-9050}
}
```

## Installation

* Clone the repository and cd to it.

```bash
git clone https://github.com/JakubSochor/BoxCars.git BoxCars
cd BoxCars
```
* (Optional, but recommended) Create virtual environment for this project - you can use **virtualenvwrapper** or following commands. **IMPORTANT NOTE:** this project is using **Python3**.

```bash
virtuenv -p /usr/bin/python3 boxcars_venv
source boxcars_venv/bin/activate
```

* Install required packages:

```bash
pip3 install -r requirements.txt 
```

* Manually download dataset https://medusa.fit.vutbr.cz/traffic/data/BoxCars116k.zip and unzip it.
* Modify `scripts/config.py` and change `BOXCARS_DATASET_ROOT` to directory where is the unzipped dataset.
* (Optional) Download trained models using `scripts/download_models.py`. To download all models to default location (`./models`) run following command (or use -h for help):

```base
python3 scripts/download_models.py --all
``` 


## Usage
### Fine-tuning of the Models
To fine-tune a model use `scripts/train_eval.py` (use -h for help). Example for ResNet50:
```bash
python3 scripts/train_eval.py --train-net ResNet50 
```
It is also possible to resume training using `--resume` argument for `train_eval.py`.

### Evaluation
The model is evaluated when the training is finished, however it is possible to evaluate saved model by running:
```bash
python3 scripts/train_eval.py --eval path-to-model.h5
```


## Trained models
We provide numerical results of models distributed with this code (use `scripts/download_models.py`). 
The processing time was measured on GTX1080 with CUDNN. The accuracy results are always shown as single image accuracy/whole track accuracy (in percents). 
We have also evaluated the method with estimated 3D bounding boxes (see paper for details) and included the results here. 
The estimated bounding boxes are in `data/estimated_3DBB.pkl`. In order to use the estimated bounding boxes, use `--estimated-3DBB path-to-pkl` argument for `train_eval.py` script.
The models which were trained with the estimated bounding boxes have suffix `_estimated3DBB`.  

Net | Original 3DBBs | Estimated 3DBBs | Image Processing Time
----|---------------:|---------------:|---------------------:
ResNet50 |  84.29/91.61 | 81.78/90.79  | 5.8ms
VGG16 | 84.10/92.09 | 81.43/90.68 | 5.4ms
VGG19 | 83.35/91.23 | 81.93/91.48  | 5.4ms
InceptionV3 | 81.51/89.86 | 79.89/89.92 | 6.1ms


## BoxCars116k dataset
The dataset was created for the paper and it is possible to download it from our [website](https://medusa.fit.vutbr.cz/traffic/data/BoxCars116k.zip)
The dataset contains 116k of images of vehicles with fine-grained labels taken from surveillance cameras under various viewpoints. 
See the paper [**BoxCars: Improving Vehicle Fine-Grained Recognition using 3D Bounding Boxes in Traffic Surveillance**](https://doi.org/10.1109/TITS.2018.2799228) for more statistics and information about dataset acquisition.
The dataset contains tracked vehicles with the same label and multiple images per track. The track is uniquely identified by its id `vehicle_id`, while each image is uniquely identified by `vehicle_id` and `instance_id`. It is possible to use class `BoxCarsDataset` from `lib/boxcars_dataset.py` for working with the dataset; however, for convenience, we describe the structure of the dataset also here. 
The dataset contains several files and folders:
* **images** - dataset images and masks 
* **atlas.pkl** - *BIG* structure with jpeg encoded images, which can be convenient as the whole structure fits the memory and it is possible to get the images on the fly. To load the atlas (or any other pkl file), you can use function `load_cache` from `lib/utils.py`. To decode the image (in RGB channel order), use the following statement.
```python
atlas = load_cache(path_to_atlas_file)
image = cv2.cvtColor(cv2.imdecode(atlas[vehicle_id][instance_id], 1), cv2.COLOR_BGR2RGB)
```

* **dataset.pkl** - contains dictionary with following fields:
```
cameras: information about used cameras (vanishing points, principal point)
samples: list of vehicles (index correspons to vehicle id). 
		 The structure contains several fields which should understandable. 
		 It also contains field instances with list of of dictionaries 
		 with information about images of the vehicle track. 
		 The flag to_camera defines whether the vehicle is going towards camera or not. 
```

* **classification_splits.pkl** - different splits (*hard*, *medium* from paper and additional *body* and *make* split). Each split contains structure `types_mapping` definig mapping from textual labels to integer labels. It also contains fields `train`, `test`, and `validation` which are lists and each element contains tuple `(vehicle_id, class_id)`.

* **verification_splits.pkl** - similar to classification splits; however, the elements in `train`, `test` are triplets `(vehicle_id1, vehicle_id2, class_id)`.

* **json_data** and **matlab_data** - converted pkl file


## Links 
* [BoxCars116k dataset](https://medusa.fit.vutbr.cz/traffic/data/BoxCars116k.zip) ([backup location](https://drive.google.com/file/d/19LHLOmmVyUS1R4ypwByfrV8KQWnz2GDT/view?usp=sharing))
* Web with our [Traffic Research](https://medusa.fit.vutbr.cz/traffic/)


================================================
FILE: data/estimated_3DBB.pkl
================================================
[File too large to display: 20.6 MB]

================================================
FILE: lib/__init__.py
================================================


================================================
FILE: lib/boxcars_data_generator.py
================================================
# -*- coding: utf-8 -*-
import cv2
import numpy as np
from keras.preprocessing.image import Iterator
from boxcars_image_transformations import alter_HSV, image_drop, unpack_3DBB, add_bb_noise_flip
import random

#%%
class BoxCarsDataGenerator(Iterator):
    def __init__(self, dataset, part, batch_size=8, training_mode=False, seed=None, generate_y = True, image_size = (224,224)):
        assert image_size == (224,224), "only images 224x224 are supported by unpack_3DBB for now, if necessary it can be changed"
        assert dataset.X[part] is not None, "load some classification split first"
        super().__init__(dataset.X[part].shape[0], batch_size, training_mode, seed)
        self.part = part
        self.generate_y = generate_y
        self.dataset = dataset
        self.image_size = image_size
        self.training_mode = training_mode
        if self.dataset.atlas is None:
            self.dataset.load_atlas()

    #%%
    def next(self):
        with self.lock:
            index_array, current_index, current_batch_size = next(self.index_generator)
        x = np.empty([current_batch_size] + list(self.image_size) + [3], dtype=np.float32)
        for i, ind in enumerate(index_array):
            vehicle_id, instance_id = self.dataset.X[self.part][ind]
            vehicle, instance, bb3d = self.dataset.get_vehicle_instance_data(vehicle_id, instance_id)
            image = self.dataset.get_image(vehicle_id, instance_id)
            if self.training_mode:
                image = alter_HSV(image) # randomly alternate color
                image = image_drop(image) # randomly remove part of the image
                bb_noise = np.clip(np.random.randn(2) * 1.5, -5, 5) # generate random bounding box movement
                flip = bool(random.getrandbits(1)) # random flip
                image, bb3d = add_bb_noise_flip(image, bb3d, flip, bb_noise) 
            image = unpack_3DBB(image, bb3d) 
            image = (image.astype(np.float32) - 116)/128.
            x[i, ...] = image
        if not self.generate_y:
            return x
        y = self.dataset.Y[self.part][index_array]
        return x, y



================================================
FILE: lib/boxcars_dataset.py
================================================
# -*- coding: utf-8 -*-
from config import BOXCARS_DATASET,BOXCARS_ATLAS,BOXCARS_CLASSIFICATION_SPLITS
from utils import load_cache
import cv2
import numpy as np

#%%
class BoxCarsDataset(object):
    def __init__(self, load_atlas = False, load_split = None, use_estimated_3DBB = False, estimated_3DBB_path = None):
        self.dataset = load_cache(BOXCARS_DATASET)
        self.use_estimated_3DBB = use_estimated_3DBB
        
        self.atlas = None
        self.split = None
        self.split_name = None
        self.estimated_3DBB = None
        self.X = {}
        self.Y = {}
        for part in ("train", "validation", "test"):
            self.X[part] = None
            self.Y[part] = None # for labels as array of 0-1 flags
            
        if load_atlas:
            self.load_atlas()
        if load_split is not None:
            self.load_classification_split(load_split)
        if self.use_estimated_3DBB:
            self.estimated_3DBB = load_cache(estimated_3DBB_path)
        
    #%%
    def load_atlas(self):
        self.atlas = load_cache(BOXCARS_ATLAS)
    
    #%%
    def load_classification_split(self, split_name):
        self.split = load_cache(BOXCARS_CLASSIFICATION_SPLITS)[split_name]
        self.split_name = split_name
       
    #%%
    def get_image(self, vehicle_id, instance_id):
        """
        returns decoded image from atlas in RGB channel order
        """
        return cv2.cvtColor(cv2.imdecode(self.atlas[vehicle_id][instance_id], 1), cv2.COLOR_BGR2RGB)
        
    #%%
    def get_vehicle_instance_data(self, vehicle_id, instance_id, original_image_coordinates=False):
        """
        original_image_coordinates: the 3DBB coordinates are in the original image space
                                    to convert them into cropped image space, it is necessary to subtract instance["3DBB_offset"]
                                    which is done if this parameter is False. 
        """
        vehicle = self.dataset["samples"][vehicle_id]
        instance = vehicle["instances"][instance_id]
        if not self.use_estimated_3DBB:
            bb3d = self.dataset["samples"][vehicle_id]["instances"][instance_id]["3DBB"]
        else:
            bb3d = self.estimated_3DBB[vehicle_id][instance_id]
            
        if not original_image_coordinates:
            bb3d = bb3d - instance["3DBB_offset"]

        return vehicle, instance, bb3d 
            
       
    #%%
    def initialize_data(self, part):
        assert self.split is not None, "load classification split first"
        assert part in self.X, "unknown part -- use: train, validation, test"
        assert self.X[part] is None, "part %s was already initialized"%part
        data = self.split[part]
        x, y = [], []
        for vehicle_id, label in data:
            num_instances = len(self.dataset["samples"][vehicle_id]["instances"])
            x.extend([(vehicle_id, instance_id) for instance_id in range(num_instances)])
            y.extend([label]*num_instances)
        self.X[part] = np.asarray(x,dtype=int)

        y = np.asarray(y,dtype=int)
        y_categorical = np.zeros((y.shape[0], self.get_number_of_classes()))
        y_categorical[np.arange(y.shape[0]), y] = 1
        self.Y[part] = y_categorical
        


    def get_number_of_classes(self):
        return len(self.split["types_mapping"])
        
        
    def evaluate(self, probabilities, part="test", top_k=1):
        samples = self.X[part]
        assert samples.shape[0] == probabilities.shape[0]
        assert self.get_number_of_classes() == probabilities.shape[1]
        part_data = self.split[part]
        probs_inds = {}
        for vehicle_id, _ in part_data:
            probs_inds[vehicle_id] = np.zeros(len(self.dataset["samples"][vehicle_id]["instances"]), dtype=int)
        for i, (vehicle_id, instance_id) in enumerate(samples):
            probs_inds[vehicle_id][instance_id] = i
            
        get_hit = lambda probs, gt: int(gt in np.argsort(probs.flatten())[-top_k:])
        hits = []
        hits_tracks = []
        for vehicle_id, label in part_data:
            inds = probs_inds[vehicle_id]
            hits_tracks.append(get_hit(np.mean(probabilities[inds, :], axis=0), label))
            for ind in inds:
                hits.append(get_hit(probabilities[ind, :], label))
                
        return np.mean(hits), np.mean(hits_tracks)
        

================================================
FILE: lib/boxcars_image_transformations.py
================================================
# -*- coding: utf-8 -*-
import cv2
import numpy as np
import random


#%%
def alter_HSV(img, change_probability = 0.6):
    if random.random() < 1-change_probability:
        return img
    addToHue = random.randint(0,179)
    addToSaturation = random.gauss(60, 20)
    addToValue = random.randint(-50,50)
    hsvVersion =  cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
    
    channels = hsvVersion.transpose(2, 0, 1)
    channels[0] = ((channels[0].astype(int) + addToHue)%180).astype(np.uint8)
    channels[1] = (np.maximum(0, np.minimum(255, (channels[1].astype(int) + addToSaturation)))).astype(np.uint8)
    channels[2] = (np.maximum(0, np.minimum(255, (channels[2].astype(int) + addToValue)))).astype(np.uint8)
    hsvVersion = channels.transpose(1,2,0)   
        
    return cv2.cvtColor(hsvVersion, cv2.COLOR_HSV2RGB)

#%%
def image_drop(img, change_probability = 0.6):
    if random.random() < 1-change_probability:
        return img
    width = random.randint(int(img.shape[1]*0.10), int(img.shape[1]*0.3))
    height = random.randint(int(img.shape[0]*0.10), int(img.shape[0]*0.3))
    x = random.randint(int(img.shape[1]*0.10), img.shape[1]-width-int(img.shape[1]*0.10))
    y = random.randint(int(img.shape[0]*0.10), img.shape[0]-height-int(img.shape[0]*0.10))
    img[y:y+height,x:x+width,:] = (np.random.rand(height,width,3)*255).astype(np.uint8)
    return img

#%%
def add_bb_noise_flip(image, bb3d, flip, bb_noise):
    bb3d = bb3d + bb_noise 
    if flip:
        bb3d[:, 0] = image.shape[1] - bb3d[:,0]
        image = cv2.flip(image, 1)
    return image, bb3d

#%%
def _unpack_side(img, origPoints, targetSize):
    origPoints = np.array(origPoints).reshape(-1,1,2)
    targetPoints = np.array([(0,0), (targetSize[0],0), (0, targetSize[1]), 
                             (targetSize[0], targetSize[1])]).reshape(-1,1,2).astype(origPoints.dtype)
    m, _ = cv2.findHomography(origPoints, targetPoints, 0)
    resultImage = cv2.warpPerspective(img, m, targetSize)
    return resultImage
    
    
#%%    
def unpack_3DBB(img, bb):
    frontal = _unpack_side(img, [bb[0], bb[1], bb[4], bb[5]], (75,124))
    side = _unpack_side(img, [bb[1], bb[2], bb[5], bb[6]], (149,124))
    roof = _unpack_side(img, [bb[0], bb[3], bb[1], bb[2]], (149,100))
    
    final = np.zeros((224,224,3), dtype=frontal.dtype)
    final[100:, 0:75] = frontal
    final[0:100, 75:] = roof
    final[100:, 75:] = side
    
    return final
    

================================================
FILE: lib/utils.py
================================================
# -*- coding: utf-8 -*-
import pickle
import os
import numpy as np
import sys

#%%
def load_cache(path, encoding="latin-1", fix_imports=True):
    """
    encoding latin-1 is default for Python2 compatibility
    """
    with open(path, "rb") as f:
        return pickle.load(f, encoding=encoding, fix_imports=True)

#%%
def save_cache(path, data):
    with open(path, "wb") as f:
        pickle.dump(data, f)

#%%
def ensure_dir(d):
    if len(d)  == 0: # for empty dirs (for compatibility with os.path.dirname("xxx.yy"))
        return
    if not os.path.exists(d):
        try:
            os.makedirs(d)
        except OSError as e:
            if e.errno != 17: # FILE EXISTS
                raise e

#%%
def parse_args(available_nets):
    import argparse
    default_cache = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "cache"))
    parser = argparse.ArgumentParser(description="BoxCars fine-grained recognition algorithm Keras re-implementation",
                                    formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--eval", type=str, default=None, help="path to model file to be evaluated")
    parser.add_argument("--resume", type=str, default=None, help="path to model file to be resumed")
    parser.add_argument("--train-net", type=str, default=available_nets[0], help="train on one of following nets: %s"%(str(available_nets)))
    parser.add_argument("--batch-size", type=int, default=8, help="batch size")
    parser.add_argument("--lr", type=float, default=0.0025, help="learning rate")
    parser.add_argument("--epochs", type=int, default=20, help="run for epochs")
    parser.add_argument("--cache", type=str, default=default_cache, help="where to store training meta-data and final model")
    parser.add_argument("--estimated-3DBB", type=str, default=None, help="use estimated 3DBBs from specified path")
    
    
    args = parser.parse_args()
    assert args.eval is None or args.resume is None, "--eval and --resume are mutually exclusive"
    if args.eval is None and args.resume is None:
        assert args.train_net in available_nets, "--train-net must be one of %s"%(str(available_nets))

    return args

 
#%%
def download_report_hook(block_num, block_size, total_size):
    downloaded = block_num*block_size
    percents = downloaded / total_size * 100
    show_str = " %.1f%%"%(percents)
    sys.stdout.write(show_str + len(show_str)*"\b")
    sys.stdout.flush()
    if downloaded >= total_size:
        print()


================================================
FILE: models/.gitignore
================================================
*
!.gitignore
!README.md


================================================
FILE: models/README.md
================================================
* Default location for downloaded models
* Use `scripts/download_models.py`

================================================
FILE: requirements.txt
================================================
appdirs==1.4.0
h5py==2.6.0
Keras==1.2.2
numpy==1.12.0
opencv-python==3.2.0.6
packaging==16.8
protobuf==3.2.0
pyparsing==2.1.10
PyYAML==3.12
scipy==0.18.1
six==1.10.0
tensorflow-gpu==1.0.0
Theano==0.8.2


================================================
FILE: scripts/_init_paths.py
================================================
# -*- coding: utf-8 -*-
import os
import sys
script_dir = os.path.dirname(__file__)
sys.path.insert(0, os.path.realpath(os.path.join(script_dir, '..', 'lib')))


================================================
FILE: scripts/config.py
================================================
# -*- coding: utf-8 -*-
import os
#%%
# change this to your location
BOXCARS_DATASET_ROOT = "/mnt/matylda1/isochor/Datasets/BoxCars116k/" 

#%%
BOXCARS_IMAGES_ROOT = os.path.join(BOXCARS_DATASET_ROOT, "images")
BOXCARS_DATASET = os.path.join(BOXCARS_DATASET_ROOT, "dataset.pkl")
BOXCARS_ATLAS = os.path.join(BOXCARS_DATASET_ROOT, "atlas.pkl")
BOXCARS_CLASSIFICATION_SPLITS = os.path.join(BOXCARS_DATASET_ROOT, "classification_splits.pkl")



================================================
FILE: scripts/download_models.py
================================================
# -*- coding: utf-8 -*-
import _init_paths
import os
import urllib.request 
import re
import argparse
import sys
from utils import ensure_dir, download_report_hook

#%%
MODELS_DIR_URL = "https://medusa.fit.vutbr.cz/traffic/data/BoxCars-models/"
SUFFIX = "h5"
DEFAULT_OUTPUT_DIR = os.path.realpath(os.path.join(os.path.dirname(__file__), "..", "models"))

#%%
with urllib.request.urlopen(MODELS_DIR_URL) as response:
    dir_listing = response.read().decode("utf-8")

model_matcher = re.compile(r'href="(.*)\.%s"'%(SUFFIX))
available_nets = model_matcher.findall(dir_listing)


#%%
parser = argparse.ArgumentParser(description="Download trained model files. Available nets: %s"%(str(available_nets)),
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument("--output-dir","-o", type=str, default=DEFAULT_OUTPUT_DIR, help="output directory where to put downloaded models")
parser.add_argument("--all", "-a", default=False, action="store_true", help="download all available models")
parser.add_argument("net_name", nargs="*")
args = parser.parse_args()

download_nets = args.net_name
if args.all:
    download_nets = available_nets
    
if len(download_nets) == 0:
    print("You need to specify nets to download or use --all to download all of them\nAVAILABLE NETS: %s\n"%(str(available_nets)))
    parser.print_usage()
    sys.exit(1)

#%%
print("Saving downloaded models to: %s"%(args.output_dir))
ensure_dir(args.output_dir)
for net in download_nets:
    if net not in available_nets:
        print("WARNING: Skipping %s because it is not available. AVAILABLE_NETS: %s"%(net, str(available_nets)))
        continue
    print("Downloading %s... "%(net), end="")
    sys.stdout.flush()
    urllib.request.urlretrieve(MODELS_DIR_URL + net + "." + SUFFIX, os.path.join(args.output_dir, "%s.%s"%(net, SUFFIX)), download_report_hook)
    


================================================
FILE: scripts/train_eval.py
================================================
# -*- coding: utf-8 -*-
import _init_paths
# this should be soon to prevent tensorflow initialization with -h parameter
from utils import ensure_dir, parse_args
args = parse_args(["ResNet50", "VGG16", "VGG19", "InceptionV3"])

# other imports
import os
import time
import sys

from boxcars_dataset import BoxCarsDataset
from boxcars_data_generator import BoxCarsDataGenerator

from keras.applications.resnet50 import ResNet50
from keras.applications.vgg16 import VGG16
from keras.applications.vgg19 import VGG19
from keras.applications.inception_v3 import InceptionV3
from keras.layers import Dense, Flatten, Dropout, AveragePooling2D
from keras.models import Model, load_model
from keras.optimizers import SGD
from keras.callbacks import ModelCheckpoint, TensorBoard


#%% initialize dataset
if args.estimated_3DBB is None:
    dataset = BoxCarsDataset(load_split="hard", load_atlas=True)
else:
    dataset = BoxCarsDataset(load_split="hard", load_atlas=True, 
                             use_estimated_3DBB = True, estimated_3DBB_path = args.estimated_3DBB)

#%% get optional path to load model
model = None
for path in [args.eval, args.resume]:
    if path is not None:
        print("Loading model from %s"%path)
        model = load_model(path)
        break

#%% construct the model as it was not passed as an argument
if model is None:
    print("Initializing new %s model ..."%args.train_net)
    if args.train_net in ("ResNet50", ):
        base_model = ResNet50(weights='imagenet', include_top=False, input_shape=(224,224,3))
        x = Flatten()(base_model.output)
        
    if args.train_net in ("VGG16", "VGG19"):
        if args.train_net == "VGG16":
            base_model = VGG16(weights='imagenet', include_top=False, input_shape=(224,224,3))
        elif args.train_net == "VGG19":
            base_model = VGG19(weights='imagenet', include_top=False, input_shape=(224,224,3))
        x = Flatten()(base_model.output)
        x = Dense(4096, activation='relu', name='fc1')(x)
        x = Dropout(0.5)(x)
        x = Dense(4096, activation='relu', name='fc2')(x)
        x = Dropout(0.5)(x)

    if args.train_net in ("InceptionV3", ):
        base_model = InceptionV3(weights='imagenet', include_top=False, input_shape=(224,224,3))
        output_dim = int(base_model.outputs[0].get_shape()[1])
        x = AveragePooling2D((output_dim, output_dim), strides=(output_dim, output_dim), name='avg_pool')(base_model.output)
        x = Flatten()(x)
            
    predictions = Dense(dataset.get_number_of_classes(), activation='softmax')(x)
    model = Model(input=base_model.input, output=predictions, name="%s%s"%(args.train_net, {True: "_estimated3DBB", False:""}[args.estimated_3DBB is not None]))
    optimizer = SGD(lr=args.lr, decay=1e-4, nesterov=True)
    model.compile(optimizer=optimizer, loss='categorical_crossentropy', metrics=["accuracy"])


print("Model name: %s"%(model.name))
if args.estimated_3DBB is not None and "estimated3DBB" not in model.name:
    print("ERROR: using estimated 3DBBs with model trained on original 3DBBs")
    sys.exit(1)
if args.estimated_3DBB is None and "estimated3DBB" in model.name:
    print("ERROR: using model trained on estimated 3DBBs and running on original 3DBBs")
    sys.exit(1)

args.output_final_model_path = os.path.join(args.cache, model.name, "final_model.h5")
args.snapshots_dir = os.path.join(args.cache, model.name, "snapshots")
args.tensorboard_dir = os.path.join(args.cache, model.name, "tensorboard")

#%% training
if args.eval is None:
    print("Training...")
    #%% initialize dataset for training
    dataset.initialize_data("train")
    dataset.initialize_data("validation")
    generator_train = BoxCarsDataGenerator(dataset, "train", args.batch_size, training_mode=True)
    generator_val = BoxCarsDataGenerator(dataset, "validation", args.batch_size, training_mode=False)


    #%% callbacks
    ensure_dir(args.tensorboard_dir)
    ensure_dir(args.snapshots_dir)
    tb_callback = TensorBoard(args.tensorboard_dir, histogram_freq=1, write_graph=False, write_images=False)
    saver_callback = ModelCheckpoint(os.path.join(args.snapshots_dir, "model_{epoch:03d}_{val_acc:.2f}.h5"), period=4 )

    #%% get initial epoch
    initial_epoch = 0
    if args.resume is not None:
        initial_epoch = int(os.path.basename(args.resume).split("_")[1]) + 1


    model.fit_generator(generator=generator_train, 
                        samples_per_epoch=generator_train.n,
                        nb_epoch=args.epochs,
                        verbose=1,
                        validation_data=generator_val,
                        nb_val_samples=generator_val.n,
                        callbacks=[tb_callback, saver_callback],
                        initial_epoch = initial_epoch,
                        )

    #%% save trained data
    print("Saving the final model to %s"%(args.output_final_model_path))
    ensure_dir(os.path.dirname(args.output_final_model_path))
    model.save(args.output_final_model_path)


#%% evaluate the model 
print("Running evaluation...")
dataset.initialize_data("test")
generator_test = BoxCarsDataGenerator(dataset, "test", args.batch_size, training_mode=False, generate_y=False)
start_time = time.time()
predictions = model.predict_generator(generator_test, generator_test.n)
end_time = time.time()
single_acc, tracks_acc = dataset.evaluate(predictions)
print(" -- Accuracy: %.2f%%"%(single_acc*100))
print(" -- Track accuracy: %.2f%%"%(tracks_acc*100))
print(" -- Image processing time: %.1fms"%((end_time-start_time)/generator_test.n*1000))