Skip to content

Commit

Permalink
feat: 🎉 implement Python Flask application with `facebook/detr-resnet…
Browse files Browse the repository at this point in the history
…-50` model for object detection
  • Loading branch information
Dan6erbond committed May 13, 2023
1 parent f9628e7 commit ee7c9a6
Show file tree
Hide file tree
Showing 5 changed files with 272 additions and 0 deletions.
160 changes: 160 additions & 0 deletions tagger/.gitignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,160 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class

# C extensions
*.so

# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST

# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec

# Installer logs
pip-log.txt
pip-delete-this-directory.txt

# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/

# Translations
*.mo
*.pot

# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal

# Flask stuff:
instance/
.webassets-cache

# Scrapy stuff:
.scrapy

# Sphinx documentation
docs/_build/

# PyBuilder
.pybuilder/
target/

# Jupyter Notebook
.ipynb_checkpoints

# IPython
profile_default/
ipython_config.py

# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version

# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock

# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock

# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/#use-with-ide
.pdm.toml

# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/

# Celery stuff
celerybeat-schedule
celerybeat.pid

# SageMath parsed files
*.sage.py

# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/

# Spyder project settings
.spyderproject
.spyproject

# Rope project settings
.ropeproject

# mkdocs documentation
/site

# mypy
.mypy_cache/
.dmypy.json
dmypy.json

# Pyre type checker
.pyre/

# pytype static type analyzer
.pytype/

# Cython debug symbols
cython_debug/

# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
#.idea/
25 changes: 25 additions & 0 deletions tagger/.vscode/launch.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
{
// Use IntelliSense to learn about possible attributes.
// Hover to view descriptions of existing attributes.
// For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
"version": "0.2.0",
"configurations": [
{
"name": "Python: Flask",
"type": "python",
"request": "launch",
"module": "flask",
"env": {
"FLASK_APP": "app.py",
"FLASK_DEBUG": "1"
},
"args": [
"run",
"--no-debugger",
"--no-reload"
],
"jinja": true,
"justMyCode": true
}
]
}
77 changes: 77 additions & 0 deletions tagger/app.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
import os
from http import HTTPStatus
from threading import Thread

import requests
import torch
from dotenv import load_dotenv
from flask import Flask
from PIL import Image
from pocketbase import PocketBase
from transformers import DetrForObjectDetection, DetrImageProcessor

load_dotenv()

app = Flask(__name__)

client = PocketBase(os.getenv("POCKETBASE_URL"))

processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")


def get_auth_token(email, password):
url = client.build_url(f"/api/admins/auth-with-password")

req = requests.post(url, json={"identity": email, "password": password})

return req.json()["token"]


def tag_file(file_id):
print(f"Tagging {file_id}")

try:
record_url = client.build_url(f"/api/collections/files/records/{file_id}")

req = requests.get(record_url)
record = req.json()

url = client.build_url(f"/api/files/files/{file_id}/{record['file']}")

image = Image.open(requests.get(url, stream=True).raw)

inputs = processor(images=image, return_tensors="pt")
outputs = model(**inputs)

# convert outputs (bounding boxes and class logits) to COCO API
# let's only keep detections with score > 0.9
target_sizes = torch.tensor([image.size[::-1]])
results = processor.post_process_object_detection(
outputs, target_sizes=target_sizes, threshold=0.9
)[0]

labels = {model.config.id2label[label.item()] for label in results["labels"]}

token = get_auth_token(
os.getenv("POCKETBASE_ADMIN_EMAIL"), os.getenv("POCKETBASE_ADMIN_PASSWORD")
)

req = requests.patch(
record_url,
json={"tagsSuggestions": list(labels)},
headers={"Authorization": f"Bearer {token}"},
)

if req.status_code != HTTPStatus.OK:
app.logger.warn(record_url, req.status_code, req.text)
except BaseException as ex:
app.logger.error(ex)


@app.route("/files/<file_id>", methods=["POST"])
def start_tag_file(file_id):
thread = Thread(target=tag_file, args=(file_id,))
thread.start()

return {"file": file_id, "model": "facebook/detr-resnet-50"}, 202
1 change: 1 addition & 0 deletions tagger/requirements-dev.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
black==23.3.0
9 changes: 9 additions & 0 deletions tagger/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
flask==2.3.2
numpy==1.24.2
pillow==9.5.0
pocketbase==0.8.1
requests==2.30.0
timm==0.6.13
torch==2.0.1
transformers==4.29.0
urllib3<2.0

0 comments on commit ee7c9a6

Please sign in to comment.