init

app.py

@@ -0,0 +1,73 @@
+from fastapi import FastAPI, UploadFile
+from ultralytics import YOLOE
+import io
+from PIL import Image
+import numpy as np
+import os
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
+import requests
+
+
+###
+
+#pip install -q "git+https://github.com/THU-MIG/yoloe.git#subdirectory=third_party/CLIP"
+#pip install -q "git+https://github.com/THU-MIG/yoloe.git#subdirectory=third_party/ml-mobileclip"
+#pip install -q "git+https://github.com/THU-MIG/yoloe.git#subdirectory=third_party/lvis-api"
+#pip install -q "git+https://github.com/THU-MIG/yoloe.git"
+
+#wget -q https://docs-assets.developer.apple.com/ml-research/datasets/mobileclip/mobileclip_blt.pt
+
+
+def init_model():
+    is_pf=True
+    model_id = "yoloe-11s"
+    # Create a models directory if it doesn't exist
+    os.makedirs("models", exist_ok=True)
+    filename = f"{model_id}-seg.pt" if not is_pf else f"{model_id}-seg-pf.pt"
+    path = hf_hub_download(repo_id="jameslahm/yoloe", filename=filename)
+    local_path = os.path.join("models", path)
+    # Download and load model
+    model = YOLOE(local_path)
+    model.eval()
+    return model
+
+app = FastAPI()
+
+# Initialize model at startup
+model = init_model()
+
+@app.post("/predict")
+async def predict(image_url: str, texts: str = "hat"):
+    # Set classes to filter
+    class_list = [text.strip() for text in texts.split(',')]
+    
+    # Download and open image from URL
+    response = requests.get(image_url)
+    image = Image.open(io.BytesIO(response.content))
+    
+    # Get text embeddings and set classes properly
+    text_embeddings = model.get_text_pe(class_list)
+    model.set_classes(class_list, text_embeddings)
+    
+    # Run inference with the PIL Image
+    results = model.predict(source=image, conf=0.25, iou=0.7)
+    
+    # Extract detection results
+    result = results[0]
+    # print(result)
+    detections = []
+    
+    for box in result.boxes:
+        detection = {
+            "class": result.names[int(box.cls[0])],
+            "confidence": float(box.conf[0]),
+            "bbox": box.xyxy[0].tolist()  # Convert bbox tensor to list
+        }
+        detections.append(detection)
+    print(detections)
+    return {"detections": detections}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

dockerfile

@@ -0,0 +1,18 @@
+# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+
+FROM python:3.10
+
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+WORKDIR /app
+
+COPY --chown=user ./requirements.txt requirements.txt
+RUN pip install --no-cache-dir --upgrade -r requirements.txt
+
+RUN wget -q https://docs-assets.developer.apple.com/ml-research/datasets/mobileclip/mobileclip_blt.pt
+
+COPY --chown=user . /app
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

requirements.txt

@@ -0,0 +1,4 @@
+"git+https://github.com/THU-MIG/yoloe.git#subdirectory=third_party/CLIP"
+"git+https://github.com/THU-MIG/yoloe.git#subdirectory=third_party/ml-mobileclip"
+"git+https://github.com/THU-MIG/yoloe.git#subdirectory=third_party/lvis-api"
+"git+https://github.com/THU-MIG/yoloe.git"