jianbing/segment-anything repository overview

FROM python:3.10-slim

WORKDIR /app

COPY sam_vit_h_4b8939.pth /data/sam_vit_h_4b8939.pth

COPY requirements.txt /app/requirements.txt
COPY demo.py /app/demo.py

RUN pip install flask 
RUN python -m pip install --upgrade pip

RUN echo 'deb https://mirrors.aliyun.com/debian/ bullseye main non-free contrib' > /etc/apt/sources.list \
  && echo 'deb-src https://mirrors.aliyun.com/debian/ bullseye main non-free contrib' >> /etc/apt/sources.list \
  && echo 'deb https://mirrors.aliyun.com/debian-security/ bullseye-security main' >> /etc/apt/sources.list \
  && echo 'deb-src https://mirrors.aliyun.com/debian-security/ bullseye-security main' >> /etc/apt/sources.list \
  && echo 'deb https://mirrors.aliyun.com/debian/ bullseye-updates main non-free contrib' >> /etc/apt/sources.list \
  && echo 'deb-src https://mirrors.aliyun.com/debian/ bullseye-updates main non-free contrib' >> /etc/apt/sources.list \
  && echo 'deb https://mirrors.aliyun.com/debian/ bullseye-backports main non-free contrib' >> /etc/apt/sources.list \
  && echo 'deb-src https://mirrors.aliyun.com/debian/ bullseye-backports main non-free contrib' >> /etc/apt/sources.list 

RUN apt update -y \
  && apt install -y git \
  && apt clean \
  && rm -rf /var/lib/apt/lists/*

RUN pip install git+https://ghproxy.com/https://github.com/facebookresearch/segment-anything.git


RUN pip3 install torch torchvision torchaudio

RUN pip install opencv-python

RUN apt update -y \
  && apt install -y ffmpeg libsm6 libxext6

EXPOSE 5000

CMD [ "python","demo.py" ]

from segment_anything import SamAutomaticMaskGenerator, sam_model_registry,SamPredictor
import numpy as np
from PIL import Image
import cv2

from flask import Flask, request,Response
from io import BytesIO

'''
pip install flask
pip install git+https://ghproxy.com/https://github.com/facebookresearch/segment-anything.git
'''

device = "cuda"
sam = sam_model_registry["vit_h"](checkpoint="/data/sam_vit_h_4b8939.pth")
sam.to(device=device)

# 初始化预测器
predictor=SamPredictor(sam)

app = Flask(__name__)

@app.route('/upload', methods=['POST'])
def upload_image():
    # 获取上传的文件
    file = request.files['image']
    # 获取文件大小
    print('获取到图片')

    # 读取图像文件
    image = cv2.imdecode(np.frombuffer(file.read(), np.uint8), cv2.IMREAD_COLOR)

    predictor.set_image(image)
    image_embedding = predictor.get_image_embedding().cpu().numpy()

    # 生成npy文件流
    with BytesIO() as buffer:
        np.save(buffer, image_embedding)
        npy_bytes = buffer.getvalue()

    # 返回npy文件流
    response = Response(npy_bytes, mimetype='application/octet-stream')
    response.headers.set('Content-Disposition', 'attachment', filename='result.npy')
    return response

if __name__ == '__main__':
    app.run(host='0.0.0.0',port=5000)