Initial commit: Recycling object detection project
- Added Fast api web application for recycling object detection - Included YOLOv8 model training notebook - Set up project structure with datasets and training directories - Added requirements.txt for dependencies - Configured .gitignore for Python virtual environment and cache files
@@ -0,0 +1,6 @@
|
||||
.venv
|
||||
__pycache__
|
||||
|
||||
augmented_dataset
|
||||
|
||||
training
|
||||
@@ -0,0 +1,54 @@
|
||||
# DS Task Recycling Project
|
||||
|
||||
This project is a toy project for training and quality assurance purposes. It involves developing a simple Flask API that processes an image (or a hardcoded image) of a motherboard and detects memory modules present on it. The API will return the image with bounding boxes drawn around each detected memory module.
|
||||
|
||||
## Project Overview
|
||||
|
||||
- **Input Types:**
|
||||
|
||||
- Image upload via the Flask API.
|
||||
- A hardcoded image for testing purposes.
|
||||
- **Dataset:**
|
||||
|
||||
- 20 pictures of motherboards with memory.
|
||||
- 20 pictures of motherboards without memory.
|
||||
- **Output:**
|
||||
|
||||
- An annotated image with bounding boxes around each detected memory module.
|
||||
For example, if there are two memory modules, two boxes are drawn; if only one is detected, then one box is drawn.
|
||||
- **Annotation Tool Suggestion:**
|
||||
|
||||
- We suggest using [makesense.ai](https://www.makesense.ai/) for manual annotation if needed.
|
||||
|
||||
## Task Details
|
||||
|
||||
The developer is required to research and answer the following questions as part of the task:
|
||||
|
||||
1. **Algorithm Choice:**
|
||||
|
||||
- Which algorithm will you use for detecting the memory modules?
|
||||
- Why do you choose this particular algorithm?
|
||||
2. **Hardware Considerations:**
|
||||
|
||||
- Does CPU or GPU have an impact on your decision? Please explain.
|
||||
3. **Video Input:**
|
||||
|
||||
- What if a video is provided instead of single images?
|
||||
- Does your approach change when processing videos? Please describe your approach.
|
||||
|
||||
## Proposed Flask API Implementation
|
||||
|
||||
1. **API Endpoints:**
|
||||
|
||||
- An endpoint for uploading images which processes and returns the annotated image.
|
||||
- An endpoint parameter for using a hardcoded image for testing purposes.
|
||||
2. **Processing Workflow:**
|
||||
|
||||
- Receive an image (either via file upload or from a hardcoded source).
|
||||
- Apply the chosen object detection algorithm to detect memory modules.
|
||||
- Draw bounding boxes around each detected memory module.
|
||||
- Return the annotated image to the user.
|
||||
|
||||
## Data Set:
|
||||
|
||||
Dataset in on the `training` folder. And there is `memory` and `no_memory` subfolder in it.
|
||||
@@ -0,0 +1,50 @@
|
||||
from fastapi import FastAPI, File, UploadFile, Request
|
||||
from fastapi.responses import HTMLResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.templating import Jinja2Templates
|
||||
import uvicorn
|
||||
from ultralytics import YOLO
|
||||
import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import io
|
||||
import os
|
||||
import base64
|
||||
|
||||
app = FastAPI()
|
||||
|
||||
# Initialize templates
|
||||
templates = Jinja2Templates(directory="templates")
|
||||
|
||||
# Load YOLO model
|
||||
current_directory = os.getcwd()
|
||||
model_path = current_directory + '\\datasets\\runs\\detect\\train8\\weights\\last.pt'
|
||||
model = YOLO(model_path)
|
||||
|
||||
@app.get("/", response_class=HTMLResponse)
|
||||
async def home(request: Request):
|
||||
return templates.TemplateResponse("index.html", {"request": request})
|
||||
|
||||
@app.post("/upload")
|
||||
async def upload_image(file: UploadFile = File(...)):
|
||||
# Read the uploaded file
|
||||
contents = await file.read()
|
||||
|
||||
# Convert to numpy array
|
||||
nparr = np.frombuffer(contents, np.uint8)
|
||||
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
||||
|
||||
# Run YOLO prediction
|
||||
results = model(img)
|
||||
|
||||
# Draw bounding boxes
|
||||
annotated_img = results[0].plot()
|
||||
|
||||
# Convert to base64
|
||||
_, img_encoded = cv2.imencode('.jpg', annotated_img)
|
||||
img_base64 = base64.b64encode(img_encoded).decode('utf-8')
|
||||
|
||||
return {"image": img_base64}
|
||||
|
||||
if __name__ == "__main__":
|
||||
uvicorn.run(app, host="localhost", port=8000)
|
||||
@@ -0,0 +1,6 @@
|
||||
|
||||
train: images/train
|
||||
val: images/val
|
||||
|
||||
nc: 1
|
||||
names: ['memory']
|
||||
|
After Width: | Height: | Size: 2.9 MiB |
|
After Width: | Height: | Size: 2.8 MiB |
|
After Width: | Height: | Size: 374 KiB |
|
After Width: | Height: | Size: 489 KiB |
|
After Width: | Height: | Size: 554 KiB |
|
After Width: | Height: | Size: 499 KiB |
|
After Width: | Height: | Size: 476 KiB |
|
After Width: | Height: | Size: 540 KiB |
|
After Width: | Height: | Size: 554 KiB |
|
After Width: | Height: | Size: 503 KiB |
|
After Width: | Height: | Size: 553 KiB |
|
After Width: | Height: | Size: 451 KiB |
|
After Width: | Height: | Size: 2.9 MiB |
|
After Width: | Height: | Size: 577 KiB |
|
After Width: | Height: | Size: 457 KiB |
|
After Width: | Height: | Size: 532 KiB |
|
After Width: | Height: | Size: 577 KiB |
|
After Width: | Height: | Size: 513 KiB |
|
After Width: | Height: | Size: 379 KiB |
|
After Width: | Height: | Size: 576 KiB |
|
After Width: | Height: | Size: 532 KiB |
|
After Width: | Height: | Size: 616 KiB |
|
After Width: | Height: | Size: 529 KiB |
|
After Width: | Height: | Size: 3.1 MiB |
|
After Width: | Height: | Size: 521 KiB |
|
After Width: | Height: | Size: 546 KiB |
|
After Width: | Height: | Size: 598 KiB |
|
After Width: | Height: | Size: 527 KiB |
|
After Width: | Height: | Size: 391 KiB |
|
After Width: | Height: | Size: 598 KiB |
|
After Width: | Height: | Size: 598 KiB |
|
After Width: | Height: | Size: 620 KiB |
|
After Width: | Height: | Size: 598 KiB |
|
After Width: | Height: | Size: 601 KiB |
|
After Width: | Height: | Size: 3.0 MiB |
|
After Width: | Height: | Size: 547 KiB |
|
After Width: | Height: | Size: 507 KiB |
|
After Width: | Height: | Size: 536 KiB |
|
After Width: | Height: | Size: 604 KiB |
|
After Width: | Height: | Size: 544 KiB |
|
After Width: | Height: | Size: 558 KiB |
|
After Width: | Height: | Size: 606 KiB |
|
After Width: | Height: | Size: 460 KiB |
|
After Width: | Height: | Size: 604 KiB |
|
After Width: | Height: | Size: 445 KiB |
|
After Width: | Height: | Size: 3.0 MiB |
|
After Width: | Height: | Size: 402 KiB |
|
After Width: | Height: | Size: 377 KiB |
|
After Width: | Height: | Size: 534 KiB |
|
After Width: | Height: | Size: 579 KiB |
|
After Width: | Height: | Size: 579 KiB |
|
After Width: | Height: | Size: 491 KiB |
|
After Width: | Height: | Size: 574 KiB |
|
After Width: | Height: | Size: 468 KiB |
|
After Width: | Height: | Size: 579 KiB |
|
After Width: | Height: | Size: 482 KiB |
|
After Width: | Height: | Size: 565 KiB |
|
After Width: | Height: | Size: 525 KiB |
|
After Width: | Height: | Size: 531 KiB |
|
After Width: | Height: | Size: 565 KiB |
|
After Width: | Height: | Size: 565 KiB |
|
After Width: | Height: | Size: 572 KiB |
|
After Width: | Height: | Size: 565 KiB |
|
After Width: | Height: | Size: 529 KiB |
|
After Width: | Height: | Size: 463 KiB |
|
After Width: | Height: | Size: 425 KiB |
|
After Width: | Height: | Size: 3.1 MiB |
|
After Width: | Height: | Size: 486 KiB |
|
After Width: | Height: | Size: 428 KiB |
|
After Width: | Height: | Size: 407 KiB |
|
After Width: | Height: | Size: 497 KiB |
|
After Width: | Height: | Size: 543 KiB |
|
After Width: | Height: | Size: 548 KiB |
|
After Width: | Height: | Size: 550 KiB |
|
After Width: | Height: | Size: 552 KiB |
|
After Width: | Height: | Size: 535 KiB |
|
After Width: | Height: | Size: 552 KiB |
|
After Width: | Height: | Size: 3.0 MiB |
|
After Width: | Height: | Size: 488 KiB |
|
After Width: | Height: | Size: 592 KiB |
|
After Width: | Height: | Size: 544 KiB |
|
After Width: | Height: | Size: 593 KiB |
|
After Width: | Height: | Size: 572 KiB |
|
After Width: | Height: | Size: 538 KiB |
|
After Width: | Height: | Size: 592 KiB |
|
After Width: | Height: | Size: 454 KiB |
|
After Width: | Height: | Size: 565 KiB |
|
After Width: | Height: | Size: 593 KiB |
|
After Width: | Height: | Size: 2.9 MiB |
|
After Width: | Height: | Size: 370 KiB |
|
After Width: | Height: | Size: 425 KiB |
|
After Width: | Height: | Size: 577 KiB |
|
After Width: | Height: | Size: 485 KiB |
|
After Width: | Height: | Size: 436 KiB |
|
After Width: | Height: | Size: 603 KiB |
|
After Width: | Height: | Size: 422 KiB |