Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

init: yolo for person detection and SAM for tracking self-annotation … #1

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions multi_object_tracking/yolo_sam/.dockerignore
Original file line number Diff line number Diff line change
@@ -0,0 +1,18 @@
# Exclude everything
_wsgi.py

# Include Dockerfile and docker-compose for reference (optional, decide based on your use case)
!Dockerfile
!docker-compose.yml

# Include Python application files
!*.py

# Include requirements files
!requirements*.txt

# Include script
!*.sh

# Exclude specific requirements if necessary
# requirements-test.txt (Uncomment if you decide to exclude this)
73 changes: 73 additions & 0 deletions multi_object_tracking/yolo_sam/Dockerfile
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime

ARG DEBIAN_FRONTEND=noninteractive
ARG TEST_ENV

WORKDIR /app

# Update Conda
RUN conda update conda -y

# Install system dependencies
RUN --mount=type=cache,target="/var/cache/apt",sharing=locked \
--mount=type=cache,target="/var/lib/apt/lists",sharing=locked \
apt-get -y update \
&& apt-get install -y git wget g++ freeglut3-dev build-essential \
libx11-dev libxmu-dev libxi-dev libglu1-mesa libglu1-mesa-dev \
libfreeimage-dev ffmpeg libsm6 libxext6 libffi-dev python3-dev \
python3-pip gcc

# Environment variables
ENV PYTHONUNBUFFERED=1 \
PYTHONDONTWRITEBYTECODE=1 \
PIP_CACHE_DIR=/.cache \
PORT=9090 \
WORKERS=2 \
THREADS=4 \
CUDA_HOME=/usr/local/cuda \
TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6+PTX;8.9;9.0" \
SEGMENT_ANYTHING_2_REPO_PATH=/segment-anything-2 \
PYTHONPATH=/app

# Install CUDA toolkit via Conda
RUN conda install -c "nvidia/label/cuda-12.1.1" cuda -y

# Install Python dependencies
COPY requirements-base.txt .
RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
pip install -r requirements-base.txt

COPY requirements.txt .
RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
pip install -r requirements.txt

# Install segment-anything-2
RUN cd / && git clone --depth 1 --branch main --single-branch https://github.com/facebookresearch/sam2.git
WORKDIR /sam2
RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
pip install -e .
RUN cd checkpoints && ./download_ckpts.sh

# Return to app working directory
WORKDIR /app

# Install test dependencies (optional)
COPY requirements-test.txt .
RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
if [ "$TEST_ENV" = "true" ]; then \
pip install -r requirements-test.txt; \
fi

# Download YOLO models
RUN /bin/sh -c 'if [ ! -f /app/models/yolov8m.pt ]; then \
yolo predict model=/app/models/yolov8m.pt source=/app/tests/car.jpg \
&& yolo predict model=/app/models/yolov8n.pt source=/app/tests/car.jpg \
&& yolo predict model=/app/models/yolov8n-cls.pt source=/app/tests/car.jpg \
&& yolo predict model=/app/models/yolov8n-seg.pt source=/app/tests/car.jpg; \
fi'

# Copy app files
COPY . ./

# Default command
CMD ["/app/start.sh"]
58 changes: 58 additions & 0 deletions multi_object_tracking/yolo_sam/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
This guide describes the simplest way to start using ML backend with Label Studio.

## Running with Docker (Recommended)

1. Start Machine Learning backend on `http://localhost:9090` with prebuilt image:

```bash
docker-compose up
```

2. Validate that backend is running

```bash
$ curl http://localhost:9090/
{"status":"UP"}
```

3. Connect to the backend from Label Studio running on the same host: go to your project `Settings -> Machine Learning -> Add Model` and specify `http://localhost:9090` as a URL.


## Building from source (Advanced)

To build the ML backend from source, you have to clone the repository and build the Docker image:

```bash
docker-compose build
```

## Running without Docker (Advanced)

To run the ML backend without Docker, you have to clone the repository and install all dependencies using pip:

```bash
python -m venv ml-backend
source ml-backend/bin/activate
pip install -r requirements.txt
```

Then you can start the ML backend:

```bash
label-studio-ml start ./dir_with_your_model
```

# Configuration
Parameters can be set in `docker-compose.yml` before running the container.


The following common parameters are available:
- `BASIC_AUTH_USER` - specify the basic auth user for the model server
- `BASIC_AUTH_PASS` - specify the basic auth password for the model server
- `LOG_LEVEL` - set the log level for the model server
- `WORKERS` - specify the number of workers for the model server
- `THREADS` - specify the number of threads for the model server

# Customization

The ML backend can be customized by adding your own models and logic inside the `./dir_with_your_model` directory.
125 changes: 125 additions & 0 deletions multi_object_tracking/yolo_sam/_wsgi.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,125 @@
import os
import argparse
import json
import logging
import logging.config

# Set a default log level if LOG_LEVEL is not defined
log_level = os.getenv("LOG_LEVEL", "INFO")

logging.config.dictConfig({
"version": 1,
"disable_existing_loggers": False, # Prevent overriding existing loggers
"formatters": {
"standard": {
"format": "[%(asctime)s] [%(levelname)s] [%(name)s::%(funcName)s::%(lineno)d] %(message)s"
}
},
"handlers": {
"console": {
"class": "logging.StreamHandler",
"level": log_level,
"stream": "ext://sys.stdout",
"formatter": "standard"
}
},
"root": {
"level": log_level,
"handlers": [
"console"
],
"propagate": True
}
})

from label_studio_ml.api import init_app
from model import NewModel


_DEFAULT_CONFIG_PATH = os.path.join(os.path.dirname(__file__), 'config.json')


def get_kwargs_from_config(config_path=_DEFAULT_CONFIG_PATH):
if not os.path.exists(config_path):
return dict()
with open(config_path) as f:
config = json.load(f)
assert isinstance(config, dict)
return config


if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Label studio')
parser.add_argument(
'-p', '--port', dest='port', type=int, default=9090,
help='Server port')
parser.add_argument(
'--host', dest='host', type=str, default='0.0.0.0',
help='Server host')
parser.add_argument(
'--kwargs', '--with', dest='kwargs', metavar='KEY=VAL', nargs='+', type=lambda kv: kv.split('='),
help='Additional LabelStudioMLBase model initialization kwargs')
parser.add_argument(
'-d', '--debug', dest='debug', action='store_true',
help='Switch debug mode')
parser.add_argument(
'--log-level', dest='log_level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], default=log_level,
help='Logging level')
parser.add_argument(
'--model-dir', dest='model_dir', default=os.path.dirname(__file__),
help='Directory where models are stored (relative to the project directory)')
parser.add_argument(
'--check', dest='check', action='store_true',
help='Validate model instance before launching server')
parser.add_argument('--basic-auth-user',
default=os.environ.get('ML_SERVER_BASIC_AUTH_USER', None),
help='Basic auth user')

parser.add_argument('--basic-auth-pass',
default=os.environ.get('ML_SERVER_BASIC_AUTH_PASS', None),
help='Basic auth pass')

args = parser.parse_args()

# setup logging level
if args.log_level:
logging.root.setLevel(args.log_level)

def isfloat(value):
try:
float(value)
return True
except ValueError:
return False

def parse_kwargs():
param = dict()
for k, v in args.kwargs:
if v.isdigit():
param[k] = int(v)
elif v == 'True' or v == 'true':
param[k] = True
elif v == 'False' or v == 'false':
param[k] = False
elif isfloat(v):
param[k] = float(v)
else:
param[k] = v
return param

kwargs = get_kwargs_from_config()

if args.kwargs:
kwargs.update(parse_kwargs())

if args.check:
print('Check "' + NewModel.__name__ + '" instance creation..')
model = NewModel(**kwargs)

app = init_app(model_class=NewModel, basic_auth_user=args.basic_auth_user, basic_auth_pass=args.basic_auth_pass)

app.run(host=args.host, port=args.port, debug=args.debug)

else:
# for uWSGI use
app = init_app(model_class=NewModel)
Empty file.
Loading