Building an Image Classification Model: Thanos vs. Joker

Introduction

As a passionate computer vision enthusiast, I embarked on an exciting journey to build an image classification model capable of distinguishing between two iconic characters: Thanos and Joker. In this article, I’ll walk you through the entire process, from data collection to model evaluation.

1. Data Collection

To create a robust dataset, I used the simple_image_download library. This Python package allowed me to download images related to both Thanos and Joker. The dataset included various poses, lighting conditions, and backgrounds, ensuring diversity for effective training.

from simple_image_download import simple_image_download as sim

response = sim.simple_image_download
response().download('thanos', 60)
response().download('joker', 60)
        

2. Model Architecture

For image classification, I designed a Convolutional Neural Network (CNN) model. Let’s break down the architecture:

model = Sequential()
model.add(Conv2D(32, (3, 3), input_shape=(64, 64, 3), activation='relu'))
model.add(MaxPooling2D(pool_size=(2, 2)))
model.add(Flatten())
model.add(Dense(units=128, activation='relu'))
model.add(Dense(units=1, activation='sigmoid'))
model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
        

• The first layer is a 2D convolutional layer with 32 filters and a 3x3 kernel size. It uses the ReLU activation function.
• Max-pooling reduces spatial dimensions.
• The flattened layer prepares the data for the fully connected layers.
• Two dense layers follow: one with 128 units (ReLU activation) and the final output layer with a single unit (sigmoid activation).

3. Data Split

I split the dataset into three subsets:

• Training set: Used for model training
• Validation set: Used for hyperparameter tuning
• Test set: Remained unseen until model evaluation

Data augmentation techniques (shear, zoom, and horizontal flip) were applied to enhance model generalization.

4. Model Training

The model was trained using the training set:

training_set = train_datagen.flow_from_directory('Dataset/train',
                                                 target_size=(64, 64),
                                                 batch_size=8,
                                                 class_mode='binary')
model.fit_generator(training_set,
                    steps_per_epoch=10,
                    epochs=50,
                    validation_data=val_set,
                    validation_steps=2)
        

5. Model Evaluation

After training, I evaluated the model’s performance on the test set:

json_file = open('model.json', 'r')
loaded_model_json = json_file.read()
json_file.close()
model = model_from_json(loaded_model_json)
model.load_weights("model.h5")

def classify(img_file):
    # Load and preprocess the test image
    # Make predictions and assign labels (Thanos or Joker)
    print(prediction, img_name)

# Iterate through test images
for f in files:
    classify(f)
        

The Result:

Github: https://github.com/heerthiraja/Deep-Learning-Projects/tree/main/Image-Classificaton-Project

Conclusion

The model successfully identified Thanos and Joker in unseen images. Feel free to explore further, fine-tune hyperparameters, and expand the dataset for even better results!

#computervision #imageclassification #deeplearning #cnn #dnn