Skip to content

Commit 8ab143d

Browse files
drduhedrduhe
committed
feat: implementing ship detection model
1 parent 6e50c55 commit 8ab143d

File tree

12 files changed

+440
-129
lines changed

12 files changed

+440
-129
lines changed

‎.gitattributes‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,5 @@
1-
assets/model_weights.pth filter=lfs diff=lfs merge=lfs -text
1+
assets/aircraft_model_weights.pth filter=lfs diff=lfs merge=lfs -text
2+
assets/ship_model_weights.pth filter=lfs diff=lfs merge=lfs -text
23
assets/images/2_planes.tiff filter=lfs diff=lfs merge=lfs -text
4+
assets/*.pth filter=lfs diff=lfs merge=lfs -text
5+
assets/images/*.tiff filter=lfs diff=lfs merge=lfs -text

‎assets/model_weights.pth‎ renamed to ‎assets/aircraft_model_weights.pth‎

File renamed without changes.

‎assets/ship_model_weights.pth‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
version https://git-lfs.github.com/spec/v1
2+
oid sha256:79818b2ccc31e7b4a91c0e286b2dcfcd1255f06ada952dd919713c7b1977ad17
3+
size 330024867

‎conda/environment-py310.yml‎

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2023-2024 Amazon.com, Inc. or its affiliates.
1+
# Copyright 2023-2025 Amazon.com, Inc. or its affiliates.
22

33
name: osml_models
44
channels:
@@ -7,12 +7,4 @@ dependencies:
77
- conda-forge::python=3.10.12
88
- conda-forge::gdal=3.7.2
99
- conda-forge::proj=9.3.0
10-
- pip:
11-
- json-logging==1.3.0
12-
- boto3==1.34.104
13-
- setuptools==68.0.0
14-
- argparse==1.4.0
15-
- flask==2.3.3
16-
- waitress==2.1.2
17-
- shapely==2.0.1
18-
- matplotlib==3.7.2
10+
- conda-forge::numpy=1.26.4

‎docker/Dockerfile‎

Lines changed: 112 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -1,62 +1,118 @@
11
# Copyright 2023-2025 Amazon.com, Inc. or its affiliates.
22

3-
# Use NVIDIA's CUDA base image
4-
FROM nvidia/cuda:11.6.2-cudnn8-devel-ubuntu18.04 AS build-env
3+
# =============================================================================
4+
# Base image: Ubuntu 22.04 + CUDA 11.6.2 (devel)
5+
# =============================================================================
6+
FROM nvidia/cuda:11.6.2-cudnn8-devel-ubuntu18.04 as osml_model
57

6-
# Set AWS to the maintainer
8+
# Set maintainer label
79
LABEL maintainer="Amazon Web Services"
810

9-
# Enable sudo access for the build session
11+
# Advertise SageMaker multi-container capability
12+
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
13+
14+
# Use root for setup
1015
USER root
1116

12-
# Update and install core build dependencies
17+
# =============================================================================
18+
# Install core build dependencies (incl. TIFF); clean apt lists in same layer
19+
# =============================================================================
1320
RUN apt-get update -y \
1421
&& apt-get upgrade -y \
15-
&& DEBIAN_FRONTEND=noninteractive apt-get install -y --fix-missing --no-install-recommends \
22+
&& DEBIAN_FRONTEND=noninteractive \
23+
apt-get install -y --fix-missing --no-install-recommends \
1624
software-properties-common build-essential ca-certificates \
1725
git make cmake wget unzip libtool automake \
1826
zlib1g-dev libsqlite3-dev pkg-config sqlite3 libcurl4-gnutls-dev \
19-
libtiff5-dev
27+
libtiff5-dev \
28+
&& rm -rf /var/lib/apt/lists/*
2029

21-
# Install Miniconda
30+
# =============================================================================
31+
# Miniconda
32+
# =============================================================================
2233
ARG MINICONDA_VERSION=Miniconda3-latest-Linux-x86_64
2334
ARG MINICONDA_URL=https://repo.anaconda.com/miniconda/${MINICONDA_VERSION}.sh
24-
ENV CONDA_TARGET_ENV=osml_model
25-
RUN wget -c ${MINICONDA_URL} \
26-
&& chmod +x ${MINICONDA_VERSION}.sh \
27-
&& ./${MINICONDA_VERSION}.sh -b -f -p /opt/conda \
28-
&& rm ${MINICONDA_VERSION}.sh \
29-
&& ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh
30-
31-
# Set our new conda target lib dirs
32-
ENV PATH=$PATH:/opt/conda/bin
33-
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/conda/lib/
34-
ENV PROJ_LIB=$PROJ_LIB:/opt/conda/share/proj
35+
ENV CONDA_DIR=/opt/conda
3536

36-
# Copy the conda environment file and create the environment
37-
COPY conda/environment-py310.yml environment.yml
38-
39-
# Accept Conda TOS before creating the environment
37+
RUN wget -c ${MINICONDA_URL} \
38+
&& chmod +x ${MINICONDA_VERSION}.sh \
39+
&& ./${MINICONDA_VERSION}.sh -b -f -p ${CONDA_DIR} \
40+
&& rm ${MINICONDA_VERSION}.sh \
41+
&& ln -s ${CONDA_DIR}/etc/profile.d/conda.sh /etc/profile.d/conda.sh
42+
43+
# Configure environment variables used by common geospatial stacks
44+
ENV CONDA_TARGET_ENV=osml_models
45+
ENV PATH=/opt/conda/envs/${CONDA_TARGET_ENV}/bin:/opt/conda/bin:$PATH
46+
ENV LD_LIBRARY_PATH=/opt/conda/envs/${CONDA_TARGET_ENV}/lib:/opt/conda/envs/${CONDA_TARGET_ENV}/lib/gdal:${LD_LIBRARY_PATH}
47+
ENV PROJ_LIB=/opt/conda/share/proj:$PROJ_LIB
48+
49+
# =============================================================================
50+
# Conda environment (py310 + GDAL/PROJ + D2 Deps)
51+
# =============================================================================
52+
COPY conda/environment-py310.yml /tmp/environment.yml
53+
54+
# Create env and minimize image size
4055
RUN conda config --set always_yes true && \
4156
conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/main && \
42-
conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r
43-
44-
RUN conda env create -n ${CONDA_TARGET_ENV} --file environment.yml && \
45-
conda clean -afy && \
57+
conda tos accept --override-channels --channel https://repo.anaconda.com/pkgs/r && \
58+
conda env create -f /tmp/environment.yml && \
4659
find /opt/conda/ -follow -type f -name '*.a' -delete && \
4760
find /opt/conda/ -follow -type f -name '*.pyc' -delete && \
4861
find /opt/conda/ -follow -type f -name '*.js.map' -delete && \
49-
rm -rf /opt/conda/pkgs
62+
rm -rf /opt/conda/pkgs && \
63+
conda clean -afy
64+
65+
# =============================================================================
66+
# Entry shim
67+
# - Ensure conda env is active for RUN/CMD/ENTRYPOINT
68+
# =============================================================================
69+
RUN cat >/entry.sh <<'BASH'
70+
#!/usr/bin/env bash
71+
set -eo pipefail
72+
73+
# Activate conda env if available
74+
if [ -f /opt/conda/etc/profile.d/conda.sh ]; then
75+
. /opt/conda/etc/profile.d/conda.sh
76+
conda activate "${CONDA_TARGET_ENV:-base}" >/dev/null 2>&1 || true
77+
fi
78+
79+
# If a command was passed, exec it; otherwise start bash
80+
if [ "$#" -gt 0 ]; then
81+
exec "$@"
82+
else
83+
exec /bin/bash
84+
fi
85+
BASH
86+
RUN chmod +x /entry.sh
87+
88+
# Make subsequent RUN use the activated env
89+
SHELL ["/entry.sh", "/bin/bash", "-c"]
90+
91+
92+
# Configure .bashrc to drop into a conda env and immediately activate our TARGET env
93+
# Note this makes python3 default to our conda managed python version
94+
RUN conda init && echo 'conda activate "${CONDA_TARGET_ENV:-base}"' >> ~/.bashrc
95+
96+
# =============================================================================
97+
# PyTorch 1.12.0 (CUDA 16.0 wheels)
98+
# =============================================================================
99+
RUN python3 -m pip install --no-cache-dir \
100+
torch==1.12.0+cu116 \
101+
torchvision==0.13.0+cu116 \
102+
-f https://download.pytorch.org/whl/torch_stable.html
50103

51-
# Activate the conda environment and install Python dependencies
52-
RUN . /opt/conda/etc/profile.d/conda.sh && conda activate ${CONDA_TARGET_ENV} && \
53-
python3 -m pip install --no-cache-dir \
104+
# =============================================================================
105+
# Detectron2 (build against target Torch/CUDA for Sagemaker Endpoints)
106+
# - Set arch list for common AWS GPUs
107+
# =============================================================================
108+
ENV FORCE_CUDA=1
109+
ARG TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing"
110+
RUN python3 -m pip install --no-cache-dir \
54111
"fvcore>=0.1.5,<0.1.6" \
55112
iopath==0.1.8 \
56113
pycocotools \
57114
omegaconf==2.1.1 \
58115
hydra-core==1.1.1 \
59-
black==21.4b2 \
60116
termcolor==1.1.0 \
61117
matplotlib==3.5.2 \
62118
yacs==0.1.8 \
@@ -65,69 +121,37 @@ RUN . /opt/conda/etc/profile.d/conda.sh && conda activate ${CONDA_TARGET_ENV} &&
65121
tqdm==4.62.3 \
66122
tensorboard==2.8.0 \
67123
opencv-contrib-python-headless==4.8.0.76 \
68-
setuptools==69.5.1
124+
setuptools==69.5.1 \
125+
'git+https://github.com/facebookresearch/detectron2.git'
69126

70-
# Install Torch with GPU support
71-
RUN . /opt/conda/etc/profile.d/conda.sh && conda activate ${CONDA_TARGET_ENV} && \
72-
python3 -m pip install --no-cache-dir \
73-
torch==1.12.0+cu116 \
74-
torchvision==0.13.0+cu116 \
75-
-f https://download.pytorch.org/whl/torch_stable.html
127+
# Final pip/conda cleanups
128+
RUN conda clean -afy && python -m pip cache purge
76129

77-
# Install Detectron2
78-
ENV FORCE_CUDA="1"
79-
ARG TORCH_CUDA_ARCH_LIST="Pascal;Volta;Turing"
80-
ENV TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}"
81-
RUN . /opt/conda/etc/profile.d/conda.sh && conda activate ${CONDA_TARGET_ENV} && \
82-
python3 -m pip install --no-cache-dir --no-deps 'git+https://github.com/facebookresearch/detectron2.git'
83-
84-
# Clean up unnecessary files
85-
RUN apt-get clean && \
86-
rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \
87-
conda clean -afy && \
88-
python -m pip cache purge
89-
90-
# Stage 2: Build the final image
91-
FROM nvidia/cuda:11.6.2-cudnn8-runtime-ubuntu18.04 AS osml_model
92-
93-
LABEL maintainer="Amazon Web Services"
94-
# Support multi-container SageMaker endpoints
95-
LABEL com.amazonaws.sagemaker.capabilities.accept-bind-to-port=true
96-
USER root
97-
98-
# Copy only the necessary files from the build environment
99-
COPY --from=build-env /opt/conda /opt/conda
100-
101-
# Set environment variables
102-
ENV CONDA_TARGET_ENV="osml_model"
103-
ENV PATH=$PATH:/opt/conda/bin
104-
ENV LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/conda/lib/
105-
ENV PROJ_LIB=$PROJ_LIB:/opt/conda/share/proj
106-
ENV PYTHONUNBUFFERED=1
107-
108-
# Set up the conda environment
109-
SHELL ["/opt/conda/bin/conda", "run", "--no-capture-output", "-n", "osml_model", "/bin/bash", "-c"]
110-
RUN echo 'conda activate "${CONDA_TARGET_ENV:-base}"' >> ~/.bashrc
111-
112-
# Copy model source and install it
113-
RUN mkdir /home/osml-models
114-
COPY . /home/osml-models
115-
116-
# Install the application dependencies
130+
# =============================================================================
131+
# Application code
132+
# =============================================================================
117133
WORKDIR /home/osml-models
118-
RUN chmod 777 --recursive .
119-
RUN python3 -m pip install --no-cache-dir .
134+
RUN mkdir -p /home/osml-models
135+
COPY . /home/osml-models
136+
RUN chmod -R 0777 . \
137+
&& python3 -m pip install --no-cache-dir .
120138

121-
# Expose the necessary ports
139+
# =============================================================================
140+
# Runtime
141+
# =============================================================================
122142
EXPOSE 8080
123143

124-
# Disable health check
144+
# Disable healthcheck (external orchestrator/SageMaker handles health)
125145
HEALTHCHECK NONE
126146

127-
# Set up a user to run the container
128-
RUN adduser --system --no-create-home --group model
129-
RUN chown -R model:model ./
147+
# Drop privileges for runtime
148+
RUN adduser --system --no-create-home --group model \
149+
&& chown -R model:model /home/osml-models \
150+
&& mkdir -p /tmp/iopath_cache && chown model:model /tmp/iopath_cache
130151
USER model
131152

132-
# Set the entry point
133-
ENTRYPOINT python3 src/aws/osml/models/$MODEL_SELECTION/app.py
153+
# Set iopath cache directory to avoid permission warnings
154+
ENV IOPATH_CACHE_DIR=/tmp/iopath_cache
155+
156+
# Expand MODEL_SELECTION, and run app
157+
ENTRYPOINT /entry.sh python /home/osml-models/src/aws/osml/models/${MODEL_SELECTION}/app.py

‎setup.cfg‎

Lines changed: 13 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -27,38 +27,29 @@ project_urls =
2727
Source = https://github.com/awslabs/osml-models
2828
Tracker = https://github.com/awslabs/osml-models/issues
2929
classifiers =
30-
Programming Language :: Python :: 3
30+
Programming Language :: Python :: 3.13
3131
Operating System :: OS Independent
32+
License :: OSI Approved :: MIT License
3233

3334
[options]
3435
zip_safe = False
35-
package_dir=
36+
package_dir =
3637
=src
37-
packages=find_namespace:
38-
python_requires = >=3.9
38+
packages = find_namespace:
39+
python_requires = >=3.10
3940
include_package_data = True
4041

4142
install_requires =
42-
json-logging==1.3.0
43-
boto3==1.34.104
44-
setuptools==68.0.0
45-
argparse==1.4.0
46-
flask==2.3.3
47-
waitress==2.1.2
48-
shapely==2.0.1
49-
matplotlib==3.7.2
43+
json-logging==1.5.1
44+
boto3==1.40.36
45+
Flask==3.1.2
46+
Waitress==3.0.2
47+
Shapely==2.1.1
48+
matplotlib==3.10.6
49+
tabulate==0.9.0
50+
yacs==0.1.8
5051

5152
[options.packages.find]
5253
where = src
5354
exclude =
5455
test
55-
56-
[options.package_data]
57-
package_data =
58-
= ["py.typed"]
59-
60-
[options.extras_require]
61-
gdal =
62-
gdal>=3.7.0
63-
test =
64-
tox

‎src/aws/osml/models/aircraft/app.py‎

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# Copyright 2023-2024 Amazon.com, Inc. or its affiliates.
1+
# Copyright 2023-2025 Amazon.com, Inc. or its affiliates.
22

33
import json
44
import os
@@ -30,6 +30,9 @@
3030
# Create our default flask app
3131
app = build_flask_app(logger)
3232

33+
# Test logging
34+
app.logger.info("Starting aircraft model application...")
35+
3336

3437
def build_predictor() -> DefaultPredictor:
3538
"""
@@ -41,7 +44,7 @@ def build_predictor() -> DefaultPredictor:
4144
# If we can't find a gpu
4245
if not torch.cuda.is_available():
4346
cfg.MODEL.DEVICE = "cpu"
44-
app.logger.warning("GPU not found, running in CPU mode!")
47+
app.logger.info("GPU not found, running in CPU mode!")
4548
# Set to only expect one class (aircraft)
4649
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1
4750
# Set the detection threshold to 90%
@@ -50,11 +53,32 @@ def build_predictor() -> DefaultPredictor:
5053
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
5154
# Path to the model weights
5255
cfg.MODEL.WEIGHTS = os.getenv(
53-
os.path.join("MODEL_WEIGHTS"), os.path.join("/home/osml-models/assets/", "model_weights.pth")
56+
os.path.join("MODEL_WEIGHTS"), os.path.join("/home/osml-models/assets/", "aircraft_model_weights.pth")
5457
)
5558

56-
# Build the detectron2 default predictor
57-
return DefaultPredictor(cfg)
59+
# Build the detectron2 default predictor with error handling for CPU mode
60+
try:
61+
# Suppress checkpoint loading warnings for expected shape mismatches
62+
import logging
63+
64+
checkpoint_logger = logging.getLogger("fvcore.common.checkpoint")
65+
original_level = checkpoint_logger.level
66+
checkpoint_logger.setLevel(logging.ERROR)
67+
68+
predictor = DefaultPredictor(cfg)
69+
70+
# Restore original logging level
71+
checkpoint_logger.setLevel(original_level)
72+
73+
return predictor
74+
except RuntimeError as e:
75+
if "NVIDIA driver" in str(e) or "CUDA" in str(e):
76+
app.logger.warning(f"CUDA error detected, forcing CPU mode: {e}")
77+
# Force CPU mode and try again
78+
cfg.MODEL.DEVICE = "cpu"
79+
return DefaultPredictor(cfg)
80+
else:
81+
raise e
5882

5983

6084
def mask_to_polygon(mask: torch.Tensor) -> List[List[float]]:
@@ -202,6 +226,7 @@ def request_to_instances(req: Request) -> Union[Instances, None]:
202226

203227
# Build our aircraft predictor
204228
aircraft_predictor = build_predictor()
229+
app.logger.info("Aircraft model predictor initialized successfully!")
205230

206231

207232
@app.route("/ping", methods=["GET"])

0 commit comments

Comments
 (0)