This is a SimpleCNN model designed to detect whether an image containing text (e.g. a scan of a document) is correctly oriented or rotated. It takes grayscale images as input, resizes them to 128x128 pixels, and outputs a prediction indicating whether the image is "Rotated" or "Normal."
Model type: Convolutional Neural Network (CNN)
Task: Binary classification (Rotated / Normal)
Input: Grayscale images (128x128)
Output: Label indicating if the image is "Rotated" or "Normal"
Framework: PyTorch
Model architecture: A simple 3-layer CNN followed by two fully connected layers
Model Description
The model consists of three convolutional layers followed by max-pooling operations. These convolutional layers extract features from the input image. The feature maps are then flattened and passed through two fully connected layers, where the final layer outputs a prediction between two classes:
Class 0 (Normal): The image is correctly oriented.
Class 1 (Rotated): The image is rotated and needs adjustment.
The model is trained on a dataset of rotated and correctly oriented grayscale images. It is capable of accurately distinguishing between the two classes and can be used in applications that involve automatic image processing or document scanning.
Usage
Inference
To use this model for inference, you can load it using Hugging Face's from_pretrained functionality and pass in an image for orientation prediction.
python
1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from safetensors.torch import load_file
5from PIL import Image
6import numpy as np
78# Define the corrected SimpleCNN architecture9classSimpleCNN(nn.Module):10def__init__(self):11super(SimpleCNN, self).__init__()12 self.conv1 = nn.Conv2d(1,16, kernel_size=3, stride=1, padding=1)# Adjusted to 16 output channels13 self.conv2 = nn.Conv2d(16,32, kernel_size=3, stride=1, padding=1)# Adjusted to 32 output channels14 self.conv3 = nn.Conv2d(32,32, kernel_size=3, stride=1, padding=1)# Adjusted to 32 output channels15 self.pool = nn.MaxPool2d(kernel_size=2, stride=2)16 self.fc1 = nn.Linear(32*16*16,32)# Adjusted input and output dimensions17 self.fc2 = nn.Linear(32,2)# Adjusted input dimension1819defforward(self, x):20 x = self.pool(F.relu(self.conv1(x)))21 x = self.pool(F.relu(self.conv2(x)))22 x = self.pool(F.relu(self.conv3(x)))23 x = x.view(x.size(0),-1)# Flatten24 x = F.relu(self.fc1(x))25 x = self.fc2(x)26return x
2728# Load the model29model = SimpleCNN()30state_dict = load_file("model.safetensors")31model.load_state_dict(state_dict)32model.eval()3334# Function to predict orientation35defpredict_orientation(image_path, model):36 img = Image.open(image_path).convert('L')# Load image in grayscale37 img = img.resize((128,128))# Resize to 128x12838 img_tensor = torch.tensor(np.array(img)/255.0, dtype=torch.float32).unsqueeze(0).unsqueeze(0)39with torch.no_grad():40 output = model(img_tensor)41 is_rotated = torch.argmax(output, dim=1).item()==142return"Rotated"if is_rotated else"Normal"4344# Example usage45result = predict_orientation("example.jpg", model)46print(f"Image Orientation: {result}")
The model performs well in scenarios where images need to be automatically detected for correct orientation. However, the performance can vary based on the image quality, input resolution, and types of rotations present in the dataset.
Limitations:
The model is trained only on 90-degree rotations, meaning performance might degrade with other types of rotations (e.g., slight tilts or partial rotations).
It is designed to work on grayscale images, so it might not perform optimally on colored or highly textured images.
Intended Use
The primary use case for this model is in scenarios where the orientation of images needs to be detected or corrected, such as:
Document scanning systems: Automatically detecting if scanned documents are oriented correctly.
Image processing pipelines: Ensuring that images are not accidentally rotated during preprocessing or ingestion.
Ethical Considerations
The model does not process or output sensitive information. However, users should be aware of potential biases that could be introduced by the training dataset (e.g., specific types of images or orientations might be overrepresented).
Citation
If you use this model, please cite the following:
bibtex
1@misc{simplecnn_orientation,
2 author = {Francesco Crescioli},
3 title = {SimpleCNN for Image Orientation Detection},
4 year = {2024},
5 howpublished = {\url{https://huggingface.co/fcrescio/rotdet}},
6}
License
This model is licensed under the Creative Commons Attribution 4.0 International (CC-BY-4.0) License.
Attribution
This model was trained using the Docmatix database, which is licensed under the MIT license. As such, the following MIT license applies to the data used in training this model:
