Upload 4 files

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.9-slim
+
+ENV PYTHONUNBUFFERED=1 \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PIP_NO_CACHE_DIR=1 \
+    MPLCONFIGDIR=/tmp/matplotlib
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y \
+    libgl1-mesa-glx \
+    libglib2.0-0 \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+
+RUN pip install --upgrade pip && \
+    pip install -r requirements.txt
+
+COPY . .
+
+EXPOSE 7860
+
+CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -0,0 +1,61 @@
+from fastapi import FastAPI
+import uvicorn
+import base64
+import cv2
+import numpy as np
+from ultralytics import YOLO
+from datetime import datetime
+from pydantic import BaseModel
+
+app = FastAPI()
+
+model = YOLO("yolo_modeln11_1502.pt")
+
+class ImageRequest(BaseModel):
+    image: str
+
+@app.get("/")
+async def root():
+    current_time = datetime.now().isoformat()
+    return {"message": "PCB Defects API works", "time": current_time}
+
+@app.post("/predict")
+async def predict(request: ImageRequest):
+    # Decode Base64
+    image_bytes = base64.b64decode(request.image)
+    np_arr = np.frombuffer(image_bytes, np.uint8)
+    image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+    if image is None:
+        return {"error": "Invalid image"}
+    
+    # Inference
+    results = model.predict(image)
+    result = results[0]
+
+    # Response
+    json_result = {}
+    class_counters = {}
+    
+    for box in result.boxes:
+        class_id = int(box.cls[0])
+        class_name = result.names[class_id]
+        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+
+        '''
+        Create unique key for each detection
+        EX: if multiple "missing_hole", keys will be:
+        "missing_hole", "missing_hole_2", ...
+        '''
+        if class_name in class_counters:
+            class_counters[class_name] += 1
+            key = f"{class_name}_{class_counters[class_name]}"
+        else:
+            class_counters[class_name] = 1
+            key = class_name
+        
+        json_result[key] = [x1, y1, x2, y2]
+    
+    return json_result
+
+if __name__ == "__main__":
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -0,0 +1,6 @@
+fastapi==0.99.1
+uvicorn==0.22.0
+pydantic==1.10.12
+numpy==1.24.2
+opencv-python==4.7.0.72
+ultralytics

yolo_modeln11_1502.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3f21751b53f8cbd3c7911a54f04b1f15e9de2443785111beef89fd5b54340ebc
+size 10770203