增加yolov5 烟火1.0检测 解决work为2的问题 解决PIL图像截断Base64剩余几个字节读不到而报错的问题
parent
4c32be5c9b
commit
11a43e5ae1
Binary file not shown.
@ -0,0 +1,45 @@
|
|||||||
|
metadata:
|
||||||
|
name: ultralytics-yolov5_firesmoke_1.0
|
||||||
|
namespace: cvat
|
||||||
|
annotations:
|
||||||
|
name: YOLO v5_firesmoke_1.0
|
||||||
|
type: detector
|
||||||
|
framework: pytorch
|
||||||
|
spec: |
|
||||||
|
[
|
||||||
|
{ "id": 0, "name": "fire" },
|
||||||
|
{ "id": 1, "name": "smoke"}
|
||||||
|
]
|
||||||
|
|
||||||
|
spec:
|
||||||
|
description: YOLO v5 via pytorch hub_firesmoke_1.0
|
||||||
|
runtime: 'python:3.6'
|
||||||
|
handler: main:handler
|
||||||
|
eventTimeout: 30s
|
||||||
|
build:
|
||||||
|
image: cvat.ultralytics-yolov5_firesmoke_1.0
|
||||||
|
baseImage: ultralytics/yolov5:latest-cpu
|
||||||
|
|
||||||
|
directives:
|
||||||
|
preCopy:
|
||||||
|
- kind: USER
|
||||||
|
value: root
|
||||||
|
- kind: RUN
|
||||||
|
value: apt update && apt install --no-install-recommends -y libglib2.0-0
|
||||||
|
- kind: WORKDIR
|
||||||
|
value: /opt/nuclio
|
||||||
|
|
||||||
|
triggers:
|
||||||
|
myHttpTrigger:
|
||||||
|
maxWorkers: 1
|
||||||
|
kind: 'http'
|
||||||
|
workerAvailabilityTimeoutMilliseconds: 10000
|
||||||
|
attributes:
|
||||||
|
maxRequestBodySize: 33554432 # 32MB
|
||||||
|
|
||||||
|
platform:
|
||||||
|
attributes:
|
||||||
|
restartPolicy:
|
||||||
|
name: always
|
||||||
|
maximumRetryCount: 3
|
||||||
|
mountMode: volume
|
||||||
@ -0,0 +1,43 @@
|
|||||||
|
import json
|
||||||
|
import base64
|
||||||
|
from PIL import Image
|
||||||
|
import io
|
||||||
|
import torch
|
||||||
|
|
||||||
|
from PIL import ImageFile
|
||||||
|
ImageFile.LOAD_TRUNCATED_IMAGES = True
|
||||||
|
|
||||||
|
def init_context(context):
|
||||||
|
context.logger.info("Init context... 0%")
|
||||||
|
|
||||||
|
# Read the DL model
|
||||||
|
model = torch.hub.load('ultralytics/yolov5', 'custom', path='/opt/nuclio/common/yolov5_firesmoke_1.0/nuclio/model.pt') # or yolov5m, yolov5l, yolov5x, custom
|
||||||
|
context.user_data.model = model
|
||||||
|
|
||||||
|
context.logger.info("Init context...100%")
|
||||||
|
|
||||||
|
def handler(context, event):
|
||||||
|
context.logger.info("Run yolo-v5_firesmoke_1.0 model")
|
||||||
|
data = event.body
|
||||||
|
buf = io.BytesIO(base64.b64decode(data["image"]))
|
||||||
|
threshold = float(data.get("threshold", 0.3))
|
||||||
|
context.user_data.model.conf = threshold
|
||||||
|
image = Image.open(buf)
|
||||||
|
yolo_results_json = context.user_data.model(image).pandas().xyxy[0].to_dict(orient='records')
|
||||||
|
|
||||||
|
encoded_results = []
|
||||||
|
for result in yolo_results_json:
|
||||||
|
encoded_results.append({
|
||||||
|
'confidence': result['confidence'],
|
||||||
|
'label': result['name'],
|
||||||
|
'points': [
|
||||||
|
result['xmin'],
|
||||||
|
result['ymin'],
|
||||||
|
result['xmax'],
|
||||||
|
result['ymax']
|
||||||
|
],
|
||||||
|
'type': 'rectangle'
|
||||||
|
})
|
||||||
|
|
||||||
|
return context.Response(body=json.dumps(encoded_results), headers={},
|
||||||
|
content_type='application/json', status_code=200)
|
||||||
Loading…
Reference in New Issue