Update README.md

README.md

@@ -1,3 +1,100 @@
----
-license: mit
----
+---
+license: mit
+datasets:
+- jxie/stl10
+---
+
+# Image Classifier
+
+This repository contains a pre-trained PyTorch model, designed for classifying images into 10 categories: airplane, bird, car, cat, deer, dog, horse, monkey, ship, and truck. The model uses a Convolutional Neural Network (CNN) architecture and can classify images based on the categories below.
+
+## Model Overview
+
+The model is a simple CNN classifier with two convolutional blocks followed by a fully connected layer. It was trained on an image dataset and can classify images into the following categories:
+
+- **0**: Airplane
+- **1**: Bird
+- **2**: Car
+- **3**: Cat
+- **4**: Deer
+- **5**: Dog
+- **6**: Horse
+- **7**: Monkey
+- **8**: Ship
+- **9**: Truck
+
+## Model Architecture
+
+The model consists of the following layers:
+1. **Conv Block 1**: Two convolutional layers with ReLU activations followed by max pooling.
+2. **Conv Block 2**: Two more convolutional layers with ReLU activations and max pooling.
+3. **Fully Connected Classifier**: A linear layer that maps the features to 10 output categories.
+
+Here’s the architecture of the model:
+```python
+class CNNV0(nn.Module):
+    def __init__(self, input_shape: int, hidden_units: int, output_shape: int):
+        super().__init__()
+        self.conv_block_1 = nn.Sequential(
+            nn.Conv2d(in_channels=input_shape, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2)
+        )
+        self.conv_block_2 = nn.Sequential(
+            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.Conv2d(in_channels=hidden_units, out_channels=hidden_units, kernel_size=3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2d(kernel_size=2)
+        )
+        self.classifier = nn.Sequential(
+            nn.Flatten(),
+            nn.Linear(in_features=hidden_units*576, out_features=output_shape)
+        )
+
+    def forward(self, x):
+        x = self.conv_block_1(x)
+        x = self.conv_block_2(x)
+        x = self.classifier(x)
+        return x
+
+## Requirements
+
+- **Python** 3.7 or higher
+- **PyTorch** 1.8 or higher
+- **torchvision** (for loading and preprocessing images)
+
+## Usage
+
+1. Clone this repository and install dependencies:
+   ```bash
+   git clone <repository-url>
+   cd <repository-folder>
+   pip install torch torchvision
+
+2. Load and use the model in your Python script:
+   ```python
+   import torch
+  from torchvision import transforms
+  from PIL import Image
+  
+  # Load the model
+  model = torch.load('model_0.pth')
+  model.eval()  # Set to evaluation mode
+  
+  # Load and preprocess the image
+  transform = transforms.Compose([
+      transforms.Resize((224, 224)),
+      transforms.ToTensor(),
+  ])
+  img = Image.open('path_to_image.jpg')
+  img = transform(img).view(1, 3, 224, 224)  # Reshape to (1, 3, 224, 224) for batch processing
+  
+  # Predict
+  with torch.no_grad():
+      output = model(img)
+      _, predicted = torch.max(output, 1)
+      print("Predicted Aircraft Type:", predicted.item())
+