Reapply "Modularize code"

This reverts commit 86bae92.

.devcontainer/devcontainer.json

@@ -0,0 +1,6 @@
+{
+  "tasks": {
+    "build": "pip install -r requirements.txt",
+    "launch": "python -m deeptuner.run"
+  }
+}

.github/workflows/publish.yml

@@ -36,4 +36,4 @@ jobs:
           pip install twine
           twine upload dist/*
       - name: Clean up
-        run: rm -rf dist build *.egg-info
+        run: rm -rf dist build *.egg-info

README.md

@@ -19,25 +19,151 @@ pip install DeepTuner
 Here is an example of how to use the package for fine-tuning models with Siamese architecture and triplet loss:
 
 ```python
-import DeepTuner
-from DeepTuner import triplet_loss, backbones, data_preprocessing, evaluation_metrics
+import os
+import json
+from sklearn.model_selection import train_test_split
+from tensorflow.keras.optimizers import Adam
+from tensorflow.keras.metrics import Mean
+from tensorflow.keras.callbacks import ReduceLROnPlateau, EarlyStopping, ModelCheckpoint
+from wandb.integration.keras import WandbMetricsLogger
+import wandb
 
-# Load and preprocess data
-data = data_preprocessing.load_data('path/to/dataset')
-triplets = data_preprocessing.create_triplets(data)
+from deeptuner.backbones.resnet import ResNetBackbone
+from deeptuner.architectures.siamese import SiameseArchitecture
+from deeptuner.losses.triplet_loss import triplet_loss
+from deeptuner.datagenerators.triplet_data_generator import TripletDataGenerator
+from deeptuner.callbacks.finetune_callback import FineTuneCallback
 
-# Initialize model backbone
-model = backbones.get_model('resnet')
+# Load configuration from JSON file
+with open('config.json', 'r') as config_file:
+    config = json.load(config_file)
 
-# Compile model with triplet loss
-model.compile(optimizer='adam', loss=triplet_loss.triplet_loss)
+data_dir = config['data_dir']
+image_size = tuple(config['image_size'])
+batch_size = config['batch_size']
+margin = config['margin']
+epochs = config['epochs']
+initial_epoch = config['initial_epoch']
+learning_rate = config['learning_rate']
+patience = config['patience']
+unfreeze_layers = config['unfreeze_layers']
 
-# Train model
-model.fit(triplets, epochs=10, batch_size=32)
+# Initialize W&B
+wandb.init(project=config['project_name'], config=config)
 
-# Evaluate model
-metrics = evaluation_metrics.evaluate_model(model, triplets)
-print(metrics)
+# Load and preprocess the data
+image_paths = []
+labels = []
+
+for label in os.listdir(data_dir):
+    label_dir = os.path.join(data_dir, label)
+    if os.path.isdir(label_dir):
+        for image_name in os.listdir(label_dir):
+            image_paths.append(os.path.join(label_dir, image_name))
+            labels.append(label)
+
+# Debugging output
+print(f"Found {len(image_paths)} images in {len(set(labels))} classes")
+
+# Split the data into training and validation sets
+train_paths, val_paths, train_labels, val_labels = train_test_split(
+    image_paths, labels, test_size=0.2, stratify=labels, random_state=42
+)
+
+# Check if the splits are non-empty
+print(f"Training on {len(train_paths)} images")
+print(f"Validating on {len(val_paths)} images")
+
+# Create data generators
+num_classes = len(set(labels))
+train_generator = TripletDataGenerator(train_paths, train_labels, batch_size, image_size, num_classes)
+val_generator = TripletDataGenerator(val_paths, val_labels, batch_size, image_size, num_classes)
+
+# Check if the generators have data
+assert len(train_generator) > 0, "Training generator is empty!"
+assert len(val_generator) > 0, "Validation generator is empty!"
+
+# Create the embedding model and freeze layers
+backbone = ResNetBackbone(input_shape=image_size + (3,))
+embedding_model = backbone.create_model()
+
+# Freeze all layers initially
+for layer in embedding_model.layers:
+    layer.trainable = False
+# Unfreeze last few layers
+for layer in embedding_model.layers[-unfreeze_layers:]:
+    layer.trainable = True
+
+# Create the siamese network
+siamese_architecture = SiameseArchitecture(input_shape=image_size + (3,), embedding_model=embedding_model)
+siamese_network = siamese_architecture.create_siamese_network()
+
+# Initialize the Siamese model
+loss_tracker = Mean(name="loss")
+siamese_model = SiameseModel(siamese_network, margin, loss_tracker)
+
+# Set up callbacks
+reduce_lr = ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=3, min_lr=1e-7, verbose=1)
+early_stopping = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True, verbose=1)
+model_checkpoint = ModelCheckpoint(
+    "models/best_siamese_model.weights.h5", 
+    save_best_only=True, 
+    save_weights_only=True, 
+    monitor='val_loss', 
+    verbose=1
+)
+embedding_checkpoint = ModelCheckpoint(
+    "models/best_embedding_model.weights.h5",
+    save_best_only=True,
+    save_weights_only=True,
+    monitor='val_loss',
+    verbose=1
+)
+fine_tune_callback = FineTuneCallback(embedding_model, patience=patience, unfreeze_layers=unfreeze_layers)
+
+# Create models directory if it doesn't exist
+os.makedirs('models', exist_ok=True)
+
+# Compile the model
+siamese_model.compile(optimizer=Adam(learning_rate=learning_rate), loss=triplet_loss(margin=margin))
+
+# Train the model
+history = siamese_model.fit(
+    train_generator,
+    validation_data=val_generator,
+    epochs=epochs,
+    initial_epoch=initial_epoch,
+    callbacks=[
+        reduce_lr, 
+        early_stopping, 
+        model_checkpoint,
+        embedding_checkpoint,
+        fine_tune_callback,
+        WandbMetricsLogger(log_freq=5)
+    ]
+)
+
+# Save the final embedding model
+embedding_model.save('models/final_embedding_model.h5')
+```
+
+### Using Configuration Files
+
+To make it easier to experiment with different hyperparameter settings, you can use a configuration file (e.g., JSON) to store hyperparameters. Here is an example of a configuration file (`config.json`):
+
+```json
+{
+    "data_dir": "path/to/your/dataset",
+    "image_size": [224, 224],
+    "batch_size": 32,
+    "margin": 1.0,
+    "epochs": 50,
+    "initial_epoch": 0,
+    "learning_rate": 0.001,
+    "patience": 5,
+    "unfreeze_layers": 10,
+    "project_name": "DeepTuner"
+}
 ```
 
-For more detailed usage and examples, please refer to the documentation.
+You can then load this configuration file in your code as shown in the usage example above.

deeptune/architectures/siamese.py

@@ -0,0 +1,19 @@
+from tensorflow.keras import layers, Model, Input
+
+class SiameseArchitecture:
+    def __init__(self, input_shape, embedding_model):
+        self.input_shape = input_shape
+        self.embedding_model = embedding_model
+
+    def create_siamese_network(self):
+        anchor_input = Input(name="anchor", shape=self.input_shape)
+        positive_input = Input(name="positive", shape=self.input_shape)
+        negative_input = Input(name="negative", shape=self.input_shape)
+
+        anchor_embedding = self.embedding_model(anchor_input)
+        positive_embedding = self.embedding_model(positive_input)
+        negative_embedding = self.embedding_model(negative_input)
+
+        outputs = [anchor_embedding, positive_embedding, negative_embedding]
+        model = Model(inputs=[anchor_input, positive_input, negative_input], outputs=outputs, name="siamese_network")
+        return model

deeptune/backbones/resnet.py

@@ -0,0 +1,27 @@
+from tensorflow.keras.applications import ResNet50
+from tensorflow.keras.layers import GlobalAveragePooling2D, Dense, Dropout, BatchNormalization
+from tensorflow.keras.models import Model
+from tensorflow.keras import Input
+
+class ResNetBackbone:
+    def __init__(self, input_shape=(224, 224, 3), weights='imagenet'):
+        self.input_shape = input_shape
+        self.weights = weights
+
+    def create_model(self):
+        inputs = Input(shape=self.input_shape)
+        base_model = ResNet50(weights=self.weights, include_top=False, input_tensor=inputs)
+        x = base_model.output
+        x = GlobalAveragePooling2D()(x)
+        x = Dense(1024, activation='relu')(x)
+        x = Dropout(0.2)(x)
+        x = BatchNormalization()(x)
+        x = Dense(512, activation='relu')(x)
+        x = Dropout(0.2)(x)
+        x = BatchNormalization()(x)
+        x = Dense(256, activation='relu')(x)
+        x = Dropout(0.2)(x)
+        x = BatchNormalization()(x)
+        outputs = Dense(128)(x)
+        model = Model(inputs, outputs, name='resnet_backbone')
+        return model

deeptune/callbacks/finetune_callback.py

@@ -0,0 +1,32 @@
+import tensorflow as tf
+from tensorflow.keras.callbacks import Callback
+from tensorflow.keras.optimizers import Adam
+
+class FineTuneCallback(Callback):
+    def __init__(self, base_model, patience=5, unfreeze_layers=10):
+        super(FineTuneCallback, self).__init__()
+        self.base_model = base_model
+        self.patience = patience
+        self.unfreeze_layers = unfreeze_layers
+        self.best_weights = None
+        self.best_loss = float('inf')
+        self.wait = 0
+
+    def on_epoch_end(self, epoch, logs=None):
+        current_loss = logs.get('val_loss')
+        if current_loss < self.best_loss:
+            self.best_loss = current_loss
+            self.best_weights = self.model.get_weights()
+            self.wait = 0
+        else:
+            self.wait += 1
+            if self.wait >= self.patience:
+                # Restore the best weights
+                self.model.set_weights(self.best_weights)
+                self.wait = 0
+                # Unfreeze the last few layers
+                for layer in self.base_model.layers[-self.unfreeze_layers:]:
+                    if hasattr(layer, 'trainable'):
+                        layer.trainable = True
+                # Recompile the model to apply the changes
+                self.model.compile(optimizer=Adam(learning_rate=1e-5))

deeptune/datagenerators/triplet_data_generator.py

@@ -0,0 +1,71 @@
+import tensorflow as tf
+from tensorflow.keras.preprocessing.image import ImageDataGenerator, load_img, img_to_array
+from sklearn.preprocessing import LabelEncoder
+import numpy as np
+from tensorflow.keras.applications import resnet50 as resnet
+
+class TripletDataGenerator(tf.keras.utils.Sequence):
+    def __init__(self, image_paths, labels, batch_size, image_size, num_classes):
+        self.image_paths = image_paths
+        self.labels = labels
+        self.batch_size = batch_size
+        self.image_size = image_size
+        self.num_classes = num_classes
+        self.label_encoder = LabelEncoder()
+        self.encoded_labels = self.label_encoder.fit_transform(labels)
+        self.image_data_generator = ImageDataGenerator(preprocessing_function=resnet.preprocess_input)
+        self.on_epoch_end()
+        print(f"Initialized TripletDataGenerator with {len(self.image_paths)} images")
+
+    def __len__(self):
+        return max(1, len(self.image_paths) // self.batch_size)  # Ensure at least one batch
+
+    def __getitem__(self, index):
+        batch_image_paths = self.image_paths[index * self.batch_size:(index + 1) * self.batch_size]
+        batch_labels = self.encoded_labels[index * self.batch_size:(index + 1) * self.batch_size]
+        return self._generate_triplet_batch(batch_image_paths, batch_labels)
+
+    def on_epoch_end(self):
+        # Shuffle the data at the end of each epoch
+        combined = list(zip(self.image_paths, self.encoded_labels))
+        np.random.shuffle(combined)
+        self.image_paths[:], self.encoded_labels[:] = zip(*combined)
+
+    def _generate_triplet_batch(self, batch_image_paths, batch_labels):
+        anchor_images = []
+        positive_images = []
+        negative_images = []
+
+        for i in range(len(batch_image_paths)):
+            anchor_path = batch_image_paths[i]
+            anchor_label = batch_labels[i]
+
+            positive_path = np.random.choice(
+                [p for p, l in zip(self.image_paths, self.encoded_labels) if l == anchor_label]
+            )
+            negative_path = np.random.choice(
+                [p for p, l in zip(self.image_paths, self.encoded_labels) if l != anchor_label]
+            )
+
+            anchor_image = load_img(anchor_path, target_size=self.image_size)
+            positive_image = load_img(positive_path, target_size=self.image_size)
+            negative_image = load_img(negative_path, target_size=self.image_size)
+
+            anchor_images.append(img_to_array(anchor_image))
+            positive_images.append(img_to_array(positive_image))
+            negative_images.append(img_to_array(negative_image))
+
+        # Convert lists to numpy arrays
+        anchor_array = np.array(anchor_images)
+        positive_array = np.array(positive_images)
+        negative_array = np.array(negative_images)
+
+        # Return inputs and dummy targets (zeros) since the loss is computed in the model
+        return (
+            {
+                "anchor": anchor_array,
+                "positive": positive_array,
+                "negative": negative_array
+            },
+            np.zeros((len(batch_image_paths),))  # Dummy target values
+        )

deeptune/losses/arcface_loss.py

@@ -0,0 +1,16 @@
+import tensorflow as tf
+from tensorflow.keras import backend as K
+
+def arcface_loss(y_true, y_pred, scale=64.0, margin=0.5):
+    y_true = tf.cast(y_true, dtype=tf.int32)
+    y_true = tf.one_hot(y_true, depth=y_pred.shape[-1])
+
+    cos_theta = y_pred
+    theta = tf.acos(K.clip(cos_theta, -1.0 + K.epsilon(), 1.0 - K.epsilon()))
+    target_logits = tf.cos(theta + margin)
+
+    logits = y_true * target_logits + (1 - y_true) * cos_theta
+    logits *= scale
+
+    loss = tf.nn.softmax_cross_entropy_with_logits(labels=y_true, logits=logits)
+    return tf.reduce_mean(loss)

deeptune/losses/triplet_loss.py

@@ -0,0 +1,11 @@
+import tensorflow as tf
+
+def triplet_loss(margin=1.0):
+    def loss(y_true, y_pred):
+        anchor, positive, negative = y_pred[0], y_pred[1], y_pred[2]
+        pos_dist = tf.reduce_sum(tf.square(anchor - positive), axis=-1)
+        neg_dist = tf.reduce_sum(tf.square(anchor - negative), axis=-1)
+        basic_loss = pos_dist - neg_dist + margin
+        loss = tf.reduce_mean(tf.maximum(basic_loss, 0.0), axis=0)
+        return loss
+    return loss