Spaces:
Running
Running
Upload 8 files
Browse files- LICENSE +21 -0
- README.md +94 -8
- requirements.txt +20 -20
- safety_classifier.py +115 -115
- safety_classifier_model.ipynb +1 -1
LICENSE
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
MIT License
|
2 |
+
|
3 |
+
Copyright (c) 2024 Darius Ardales
|
4 |
+
|
5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
6 |
+
of this software and associated documentation files (the "Software"), to deal
|
7 |
+
in the Software without restriction, including without limitation the rights
|
8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
9 |
+
copies of the Software, and to permit persons to whom the Software is
|
10 |
+
furnished to do so, subject to the following conditions:
|
11 |
+
|
12 |
+
The above copyright notice and this permission notice shall be included in all
|
13 |
+
copies or substantial portions of the Software.
|
14 |
+
|
15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
21 |
+
SOFTWARE.
|
README.md
CHANGED
@@ -1,11 +1,97 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
---
|
2 |
-
|
3 |
-
|
4 |
-
|
5 |
-
|
6 |
-
|
7 |
-
|
8 |
-
|
|
|
|
|
|
|
|
|
9 |
---
|
10 |
|
11 |
-
|
|
|
1 |
+
# Towards a Safer Construction Environment: Evaluating a Simple CNN for Safety Classification
|
2 |
+
### By: Darius Vincent C. Ardales
|
3 |
+
|
4 |
+
This project explores workplace safety compliance in construction sites by evaluating a Convolutional Neural Network (CNN) for binary classification of images into "safe" (individuals wearing safety equipment) and "unsafe" (individuals not wearing safety equipment). It includes training, testing, and deploying the model for real-time classification through a FastAPI app, containerized with Docker, and hosted on Hugging Face Spaces.
|
5 |
+
|
6 |
+
## Prerequisites
|
7 |
+
|
8 |
+
Before you begin, ensure you have the following:
|
9 |
+
|
10 |
+
- **pip** (preferred if you also have `conda` for environment management)
|
11 |
+
- **Python 3.9.18** (strictly this version for compatibility)
|
12 |
+
- A virtual environment (optional but highly recommended)
|
13 |
+
- **Operating System**: This project was developed on a Microsoft Windows 11 laptop.
|
14 |
+
|
15 |
+
---
|
16 |
+
|
17 |
+
## Project Structure
|
18 |
+
|
19 |
+
Below is an organized overview of the project's folders, files, and their roles:
|
20 |
+
|
21 |
+
### **Folders**
|
22 |
+
- **`data`**
|
23 |
+
- Contains two subfolders:
|
24 |
+
- `safe`: Original images classified as "safe".
|
25 |
+
- `unsafe`: Original images classified as "unsafe".
|
26 |
+
- **`augmented`**
|
27 |
+
- Contains two subfolders:
|
28 |
+
- `safe`: Augmented "safe" images specifically used for training.
|
29 |
+
- `unsafe`: Augmented "unsafe" images specifically used for training.
|
30 |
+
|
31 |
+
### **Files**
|
32 |
+
- **`safety_classifier_model.ipynb`**
|
33 |
+
- Jupyter notebook where all steps of the project are executed, including preprocessing, augmentation, model creation, training, testing, and analysis.
|
34 |
+
- **`safety_classifier.py`**
|
35 |
+
- Executable FastAPI code for real-time binary image classification (safe/unsafe) using the trained SafetyCNN model.
|
36 |
+
- **`Dockerfile`**
|
37 |
+
- Containerizes the FastAPI app for deployment on Hugging Face Spaces or other hosting platforms.
|
38 |
+
- **`safety_model.pth`**
|
39 |
+
- Saved PyTorch model of the trained CNN for safety classification.
|
40 |
+
- **`requirements.txt`**
|
41 |
+
- List of dependencies needed to run `safety_classifier.py`.
|
42 |
+
|
43 |
+
---
|
44 |
+
|
45 |
+
## Steps to Run the Project
|
46 |
+
|
47 |
+
Follow these steps to clone, set up, and run the project locally:
|
48 |
+
|
49 |
+
### 1. **Clone the Repository**
|
50 |
+
```bash
|
51 |
+
git clone https://github.com/riu-rd/Worksite-Safety-Monitoring.git
|
52 |
+
cd Worksite-Safety-Monitoring
|
53 |
+
```
|
54 |
+
|
55 |
+
### 2. Download the Model (Needed step because of GitHub memory constraints)
|
56 |
+
- Download safety_model.pth from the author's Hugging Face [HERE](https://huggingface.co/spaces/riu-rd/safety_classifier/blob/main/safety_model.pth).
|
57 |
+
- Place the file in the same directory as the project files.
|
58 |
+
|
59 |
+
### 3. Set Up a Virtual Environment
|
60 |
+
- If you don’t have a virtual environment, create one using `pip` or `conda`. Make sure the Python version is `3.9.18`.
|
61 |
+
- Example using `pip`:
|
62 |
+
```bash
|
63 |
+
python -m venv env
|
64 |
+
source env/bin/activate # On Linux/MacOS
|
65 |
+
env\Scripts\activate # On Windows
|
66 |
+
```
|
67 |
+
- Example using `conda`:
|
68 |
+
```bash
|
69 |
+
conda create --name safety_env python=3.9.18
|
70 |
+
conda activate safety_env
|
71 |
+
```
|
72 |
+
|
73 |
+
### 4. Install Dependencies
|
74 |
+
```bash
|
75 |
+
pip install -r requirements.txt
|
76 |
+
```
|
77 |
+
|
78 |
+
### 5. Run the FastAPI App Locally
|
79 |
+
```bash
|
80 |
+
uvicorn safety_classifier:app --reload
|
81 |
+
```
|
82 |
+
|
83 |
---
|
84 |
+
|
85 |
+
## Deployment
|
86 |
+
|
87 |
+
- The FastAPI app is containerized using Docker and hosted live on Hugging Face Spaces. To view the project live, visit: [Live Link](https://riu-rd-safety-classifier.hf.space/)
|
88 |
+
|
89 |
+
---
|
90 |
+
|
91 |
+
## Notes and Recommendations
|
92 |
+
- Ensure that your environment matches the prerequisites for a smooth setup.
|
93 |
+
- For issues or contributions, please open an issue in the repository or submit a pull request.
|
94 |
+
|
95 |
---
|
96 |
|
97 |
+
# Happy Worksite Safety Classifying!
|
requirements.txt
CHANGED
@@ -1,20 +1,20 @@
|
|
1 |
-
# Python version
|
2 |
-
# python==3.9.18
|
3 |
-
|
4 |
-
# FastAPI and related dependencies
|
5 |
-
fastapi==0.95.1
|
6 |
-
uvicorn==0.22.0
|
7 |
-
|
8 |
-
# Torch and torchvision
|
9 |
-
torch==1.13.1
|
10 |
-
torchvision==0.14.1
|
11 |
-
typing-extensions==4.6.3 # Sometimes required for compatibility with PyTorch
|
12 |
-
|
13 |
-
# Image processing
|
14 |
-
Pillow==9.4.0
|
15 |
-
opencv-python-headless==4.7.0.72
|
16 |
-
|
17 |
-
# Other dependencies
|
18 |
-
numpy==1.23.5
|
19 |
-
matplotlib==3.6.3
|
20 |
-
python-multipart
|
|
|
1 |
+
# Python version
|
2 |
+
# python==3.9.18
|
3 |
+
|
4 |
+
# FastAPI and related dependencies
|
5 |
+
fastapi==0.95.1
|
6 |
+
uvicorn==0.22.0
|
7 |
+
|
8 |
+
# Torch and torchvision
|
9 |
+
torch==1.13.1
|
10 |
+
torchvision==0.14.1
|
11 |
+
typing-extensions==4.6.3 # Sometimes required for compatibility with PyTorch
|
12 |
+
|
13 |
+
# Image processing
|
14 |
+
Pillow==9.4.0
|
15 |
+
opencv-python-headless==4.7.0.72
|
16 |
+
|
17 |
+
# Other dependencies
|
18 |
+
numpy==1.23.5
|
19 |
+
matplotlib==3.6.3
|
20 |
+
python-multipart
|
safety_classifier.py
CHANGED
@@ -1,115 +1,115 @@
|
|
1 |
-
from fastapi import FastAPI, File, UploadFile, HTTPException
|
2 |
-
from fastapi.responses import JSONResponse, RedirectResponse
|
3 |
-
import torch
|
4 |
-
import torch.nn as nn
|
5 |
-
import torch.nn.functional as F
|
6 |
-
from torchvision import transforms
|
7 |
-
from PIL import Image
|
8 |
-
import os
|
9 |
-
import io
|
10 |
-
import numpy as np
|
11 |
-
import matplotlib.pyplot as plt
|
12 |
-
import cv2 as cv
|
13 |
-
from typing import List
|
14 |
-
|
15 |
-
# Create FastAPI app
|
16 |
-
app = FastAPI()
|
17 |
-
|
18 |
-
# Load the pre-trained SafetyCNN model
|
19 |
-
class SafetyCNN(nn.Module):
|
20 |
-
def __init__(self):
|
21 |
-
super().__init__()
|
22 |
-
self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=5, stride=1)
|
23 |
-
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
|
24 |
-
self.conv2 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=5, stride=1)
|
25 |
-
self.dropout = nn.Dropout(0.5)
|
26 |
-
self.fc1 = nn.Linear(157 * 157 * 24, 64)
|
27 |
-
self.fc2 = nn.Linear(64, 1)
|
28 |
-
|
29 |
-
def forward(self, x):
|
30 |
-
x = self.pool(F.relu(self.conv1(x)))
|
31 |
-
x = self.pool(F.relu(self.conv2(x)))
|
32 |
-
x = torch.flatten(x, 1)
|
33 |
-
x = F.relu(self.fc1(x))
|
34 |
-
x = self.dropout(x)
|
35 |
-
x = self.fc2(x)
|
36 |
-
return x
|
37 |
-
|
38 |
-
# Initialize model and load pre-trained weights
|
39 |
-
safety_model = SafetyCNN()
|
40 |
-
safety_model.load_state_dict(torch.load("safety_model.pth", map_location=torch.device('cpu')))
|
41 |
-
safety_model.eval()
|
42 |
-
|
43 |
-
# Define transformations
|
44 |
-
tta_transforms = [
|
45 |
-
transforms.Compose([]),
|
46 |
-
transforms.Compose([transforms.RandomHorizontalFlip(p=1.0)]),
|
47 |
-
transforms.Compose([transforms.RandomRotation(degrees=30)]),
|
48 |
-
transforms.Compose([transforms.RandomResizedCrop(size=(640, 640), scale=(0.8, 1.0))])
|
49 |
-
]
|
50 |
-
|
51 |
-
# Utility function to classify image with TTA
|
52 |
-
def classify_image_with_tta(image: Image.Image, model, tta_transforms: List[transforms.Compose], num_tta=4):
|
53 |
-
# List to accumulate predictions from TTA versions
|
54 |
-
augmented_predictions = []
|
55 |
-
|
56 |
-
# Apply each TTA transformation to the image
|
57 |
-
for i in range(num_tta):
|
58 |
-
tta_transform = tta_transforms[i % len(tta_transforms)]
|
59 |
-
augmented_image = tta_transform(image)
|
60 |
-
|
61 |
-
# Preprocess the image for the model
|
62 |
-
transform = transforms.Compose([
|
63 |
-
transforms.Resize((640, 640)),
|
64 |
-
transforms.ToTensor(),
|
65 |
-
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
66 |
-
])
|
67 |
-
input_tensor = transform(augmented_image).unsqueeze(0) # Add batch dimension (1, 3, 640, 640)
|
68 |
-
|
69 |
-
# Run the model on the augmented image
|
70 |
-
with torch.no_grad():
|
71 |
-
output = model(input_tensor).squeeze(1)
|
72 |
-
|
73 |
-
# Apply sigmoid to get the probability
|
74 |
-
prob = torch.sigmoid(output).item()
|
75 |
-
augmented_predictions.append(prob)
|
76 |
-
|
77 |
-
# Average predictions over all TTA versions
|
78 |
-
avg_prob = np.mean(augmented_predictions)
|
79 |
-
|
80 |
-
# Set a threshold of 0.5 for binary classification
|
81 |
-
prediction = 1 if avg_prob > 0.5 else 0
|
82 |
-
|
83 |
-
return prediction, avg_prob
|
84 |
-
|
85 |
-
@app.get("/")
|
86 |
-
async def docs():
|
87 |
-
return RedirectResponse(url="/docs")
|
88 |
-
|
89 |
-
# FastAPI endpoint to upload an image and classify it
|
90 |
-
@app.post("/classify")
|
91 |
-
async def classify_image(file: UploadFile = File(...)):
|
92 |
-
try:
|
93 |
-
# Read the uploaded file
|
94 |
-
contents = await file.read()
|
95 |
-
image = Image.open(io.BytesIO(contents)).convert("RGB")
|
96 |
-
except Exception as e:
|
97 |
-
raise HTTPException(status_code=400, detail="Invalid image file")
|
98 |
-
|
99 |
-
# Classify the image using the model
|
100 |
-
prediction, avg_prob = classify_image_with_tta(image, safety_model, tta_transforms, num_tta=4)
|
101 |
-
|
102 |
-
# Create a response message
|
103 |
-
result = {
|
104 |
-
"prediction": "Unsafe" if prediction == 1 else "Safe",
|
105 |
-
"probability": avg_prob
|
106 |
-
}
|
107 |
-
return JSONResponse(content=result)
|
108 |
-
|
109 |
-
# Optional: Run with uvicorn if needed
|
110 |
-
# if __name__ == "__main__":
|
111 |
-
# import uvicorn
|
112 |
-
# uvicorn.run(app, host="0.0.0.0", port=8000)
|
113 |
-
|
114 |
-
# To run the server:
|
115 |
-
# uvicorn safety_classifier:app --reload
|
|
|
1 |
+
from fastapi import FastAPI, File, UploadFile, HTTPException
|
2 |
+
from fastapi.responses import JSONResponse, RedirectResponse
|
3 |
+
import torch
|
4 |
+
import torch.nn as nn
|
5 |
+
import torch.nn.functional as F
|
6 |
+
from torchvision import transforms
|
7 |
+
from PIL import Image
|
8 |
+
import os
|
9 |
+
import io
|
10 |
+
import numpy as np
|
11 |
+
import matplotlib.pyplot as plt
|
12 |
+
import cv2 as cv
|
13 |
+
from typing import List
|
14 |
+
|
15 |
+
# Create FastAPI app
|
16 |
+
app = FastAPI()
|
17 |
+
|
18 |
+
# Load the pre-trained SafetyCNN model
|
19 |
+
class SafetyCNN(nn.Module):
|
20 |
+
def __init__(self):
|
21 |
+
super().__init__()
|
22 |
+
self.conv1 = nn.Conv2d(in_channels=3, out_channels=12, kernel_size=5, stride=1)
|
23 |
+
self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
|
24 |
+
self.conv2 = nn.Conv2d(in_channels=12, out_channels=24, kernel_size=5, stride=1)
|
25 |
+
self.dropout = nn.Dropout(0.5)
|
26 |
+
self.fc1 = nn.Linear(157 * 157 * 24, 64)
|
27 |
+
self.fc2 = nn.Linear(64, 1)
|
28 |
+
|
29 |
+
def forward(self, x):
|
30 |
+
x = self.pool(F.relu(self.conv1(x)))
|
31 |
+
x = self.pool(F.relu(self.conv2(x)))
|
32 |
+
x = torch.flatten(x, 1)
|
33 |
+
x = F.relu(self.fc1(x))
|
34 |
+
x = self.dropout(x)
|
35 |
+
x = self.fc2(x)
|
36 |
+
return x
|
37 |
+
|
38 |
+
# Initialize model and load pre-trained weights
|
39 |
+
safety_model = SafetyCNN()
|
40 |
+
safety_model.load_state_dict(torch.load("safety_model.pth", map_location=torch.device('cpu')))
|
41 |
+
safety_model.eval()
|
42 |
+
|
43 |
+
# Define transformations
|
44 |
+
tta_transforms = [
|
45 |
+
transforms.Compose([]),
|
46 |
+
transforms.Compose([transforms.RandomHorizontalFlip(p=1.0)]),
|
47 |
+
transforms.Compose([transforms.RandomRotation(degrees=30)]),
|
48 |
+
transforms.Compose([transforms.RandomResizedCrop(size=(640, 640), scale=(0.8, 1.0))])
|
49 |
+
]
|
50 |
+
|
51 |
+
# Utility function to classify image with TTA
|
52 |
+
def classify_image_with_tta(image: Image.Image, model, tta_transforms: List[transforms.Compose], num_tta=4):
|
53 |
+
# List to accumulate predictions from TTA versions
|
54 |
+
augmented_predictions = []
|
55 |
+
|
56 |
+
# Apply each TTA transformation to the image
|
57 |
+
for i in range(num_tta):
|
58 |
+
tta_transform = tta_transforms[i % len(tta_transforms)]
|
59 |
+
augmented_image = tta_transform(image)
|
60 |
+
|
61 |
+
# Preprocess the image for the model
|
62 |
+
transform = transforms.Compose([
|
63 |
+
transforms.Resize((640, 640)),
|
64 |
+
transforms.ToTensor(),
|
65 |
+
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
|
66 |
+
])
|
67 |
+
input_tensor = transform(augmented_image).unsqueeze(0) # Add batch dimension (1, 3, 640, 640)
|
68 |
+
|
69 |
+
# Run the model on the augmented image
|
70 |
+
with torch.no_grad():
|
71 |
+
output = model(input_tensor).squeeze(1)
|
72 |
+
|
73 |
+
# Apply sigmoid to get the probability
|
74 |
+
prob = torch.sigmoid(output).item()
|
75 |
+
augmented_predictions.append(prob)
|
76 |
+
|
77 |
+
# Average predictions over all TTA versions
|
78 |
+
avg_prob = np.mean(augmented_predictions)
|
79 |
+
|
80 |
+
# Set a threshold of 0.5 for binary classification
|
81 |
+
prediction = 1 if avg_prob > 0.5 else 0
|
82 |
+
|
83 |
+
return prediction, avg_prob
|
84 |
+
|
85 |
+
@app.get("/")
|
86 |
+
async def docs():
|
87 |
+
return RedirectResponse(url="/docs")
|
88 |
+
|
89 |
+
# FastAPI endpoint to upload an image and classify it
|
90 |
+
@app.post("/classify")
|
91 |
+
async def classify_image(file: UploadFile = File(...)):
|
92 |
+
try:
|
93 |
+
# Read the uploaded file
|
94 |
+
contents = await file.read()
|
95 |
+
image = Image.open(io.BytesIO(contents)).convert("RGB")
|
96 |
+
except Exception as e:
|
97 |
+
raise HTTPException(status_code=400, detail="Invalid image file")
|
98 |
+
|
99 |
+
# Classify the image using the model
|
100 |
+
prediction, avg_prob = classify_image_with_tta(image, safety_model, tta_transforms, num_tta=4)
|
101 |
+
|
102 |
+
# Create a response message
|
103 |
+
result = {
|
104 |
+
"prediction": "Unsafe" if prediction == 1 else "Safe",
|
105 |
+
"probability": avg_prob
|
106 |
+
}
|
107 |
+
return JSONResponse(content=result)
|
108 |
+
|
109 |
+
# Optional: Run with uvicorn if needed
|
110 |
+
# if __name__ == "__main__":
|
111 |
+
# import uvicorn
|
112 |
+
# uvicorn.run(app, host="0.0.0.0", port=8000)
|
113 |
+
|
114 |
+
# To run the server:
|
115 |
+
# uvicorn safety_classifier:app --reload
|
safety_classifier_model.ipynb
CHANGED
@@ -1113,7 +1113,7 @@
|
|
1113 |
"cell_type": "markdown",
|
1114 |
"metadata": {},
|
1115 |
"source": [
|
1116 |
-
"# Model Evaluation (
|
1117 |
]
|
1118 |
},
|
1119 |
{
|
|
|
1113 |
"cell_type": "markdown",
|
1114 |
"metadata": {},
|
1115 |
"source": [
|
1116 |
+
"# Model Evaluation (Test-Time Augmentation)"
|
1117 |
]
|
1118 |
},
|
1119 |
{
|