Move jetson deployment code
This commit is contained in:
@@ -0,0 +1,165 @@
|
||||
import argparse
|
||||
import cv2
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
|
||||
|
||||
def load_models(yolo_path: str, resnet_path: str):
|
||||
"""Load the models for two-stage classification.
|
||||
|
||||
:param str yolo_path: path to yolo weights
|
||||
:param str resnet_path: path to resnet weights
|
||||
:returns: tuple of models
|
||||
|
||||
"""
|
||||
first_stage = torch.hub.load("WongKinYiu/yolov7",
|
||||
"custom",
|
||||
yolo_path,
|
||||
trust_repo=True)
|
||||
second_stage = torch.load(resnet_path)
|
||||
return (first_stage, second_stage)
|
||||
|
||||
|
||||
def detect(img_path: str, yolo_path: str, resnet_path: str):
|
||||
"""Load an image, detect individual plants and label them as
|
||||
healthy or wilted.
|
||||
|
||||
:param str img_path: path to image
|
||||
:param str yolo_path: path to yolo weights
|
||||
:param str resnet_path: path to resnet weights
|
||||
:returns: tuple of recent image and dict of bounding boxes and
|
||||
their predictions
|
||||
|
||||
"""
|
||||
img = cv2.imread(img_path)
|
||||
original = img.copy()
|
||||
(first_stage, second_stage) = load_models(yolo_path, resnet_path)
|
||||
|
||||
# Get bounding boxes from object detection model
|
||||
box_coords = get_boxes(first_stage, img)
|
||||
|
||||
box_coords.sort_values(by=['xmin'], ignore_index=True, inplace=True)
|
||||
|
||||
predictions = {}
|
||||
for idx, row in box_coords.iterrows():
|
||||
xmin, xmax = int(row['xmin']), int(row['xmax'])
|
||||
ymin, ymax = int(row['ymin']), int(row['ymax'])
|
||||
|
||||
# Get tensor of ROI in BGR
|
||||
cropped_image = get_cutout(img.copy(), xmin, xmax, ymin, ymax)
|
||||
|
||||
# Classify ROI in RGB
|
||||
predictions[idx] = classify(second_stage, cropped_image[..., ::-1])
|
||||
|
||||
# Draw bounding box and number on original image
|
||||
original = cv2.rectangle(original, (xmin, ymin), (xmax, ymax),
|
||||
(0, 255, 0), 2)
|
||||
original = cv2.putText(original, str(idx), (xmin + 5, ymin + 25),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 0), 4,
|
||||
cv2.LINE_AA)
|
||||
original = cv2.putText(original, str(idx), (xmin + 5, ymin + 25),
|
||||
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255),
|
||||
2, cv2.LINE_AA)
|
||||
|
||||
return (original, predictions)
|
||||
|
||||
|
||||
def get_boxes(model, img):
|
||||
"""Run object detection model on an image and get the bounding box
|
||||
coordinates of all matches.
|
||||
|
||||
:param model: object detection model (YOLO)
|
||||
:param img: opencv2 image object
|
||||
:returns: pandas dataframe of matches
|
||||
|
||||
"""
|
||||
with torch.no_grad():
|
||||
box_coords = model(img[..., ::-1], size=640)
|
||||
return box_coords.pandas().xyxy[0]
|
||||
|
||||
|
||||
def classify(model, img):
|
||||
"""Classify img with object classification model.
|
||||
|
||||
:param model: object classification model
|
||||
:param img: opencv2 image object in RGB
|
||||
:returns: tensor of class predictions
|
||||
|
||||
"""
|
||||
# Transform image for ResNet
|
||||
data_transforms = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
|
||||
])
|
||||
|
||||
img = data_transforms(img.copy())
|
||||
with torch.no_grad():
|
||||
out = model(img.unsqueeze(0))
|
||||
# Apply softmax to get percentage confidence of classes
|
||||
out = torch.nn.functional.softmax(out, dim=1)[0] * 100
|
||||
return out
|
||||
|
||||
|
||||
def get_cutout(img, xmin, xmax, ymin, ymax):
|
||||
"""Cut out a bounding box from an image and transform it for
|
||||
object classification model.
|
||||
|
||||
:param img: opencv2 image object in BGR
|
||||
:param int xmin: start of bounding box on x axis
|
||||
:param int xmax: end of bounding box on x axis
|
||||
:param int ymin: start of bounding box on y axis
|
||||
:param int ymax: end of bounding box on y axis
|
||||
:returns: tensor of cropped image in BGR
|
||||
|
||||
"""
|
||||
cropped_image = img[ymin:ymax, xmin:xmax]
|
||||
return cropped_image
|
||||
|
||||
|
||||
def export_to_onnx(yolo_path: str, resnet_path: str):
|
||||
"""Export the models to onnx.
|
||||
|
||||
:param yolo_path: path to yolo weights
|
||||
:param resnet_path: path to resnet weights
|
||||
:returns: None
|
||||
|
||||
"""
|
||||
(first, second) = load_models(yolo_path, resnet_path)
|
||||
first.eval()
|
||||
second.eval()
|
||||
|
||||
first_x = torch.randn((1, 3, 640, 640), requires_grad=True)
|
||||
second_x = torch.randn((1, 3, 224, 224), requires_grad=True)
|
||||
|
||||
torch.onnx.export(first,
|
||||
first_x,
|
||||
'yolo.onnx',
|
||||
export_params=True,
|
||||
do_constant_folding=True,
|
||||
input_names=['input'],
|
||||
output_names=['output'])
|
||||
|
||||
torch.onnx.export(second,
|
||||
second_x,
|
||||
'resnet.onnx',
|
||||
export_params=True,
|
||||
do_constant_folding=True,
|
||||
input_names=['input'],
|
||||
output_names=['output'])
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--source', type=str, help='image file or webcam')
|
||||
parser.add_argument('--onnx',
|
||||
action='store_true',
|
||||
dest='onnx',
|
||||
help='export models to onnx')
|
||||
opt = parser.parse_args()
|
||||
|
||||
if opt.source:
|
||||
detect(opt.source, 'runs/train/yolov7-custom7/weights/best.pt',
|
||||
'resnet.pt')
|
||||
if opt.onnx:
|
||||
export_to_onnx('runs/train/yolov7-custom7/weights/best.pt',
|
||||
'resnet.pt')
|
||||
@@ -0,0 +1,52 @@
|
||||
import atexit
|
||||
import jetson.utils
|
||||
|
||||
from flask import Flask
|
||||
from apscheduler.schedulers.background import BackgroundScheduler
|
||||
|
||||
from multiprocessing import Manager
|
||||
|
||||
from model import detect
|
||||
|
||||
app = Flask(__name__)
|
||||
scheduler = BackgroundScheduler(daemon=True)
|
||||
manager = Manager()
|
||||
pred = manager.dict()
|
||||
|
||||
|
||||
@scheduler.task('interval', id='get_pred', minutes=30, misfire_grace_time=900)
|
||||
def get_pred():
|
||||
img = take_image('./current_image.jpg')
|
||||
print('Job 1 executed')
|
||||
|
||||
|
||||
def take_image(img_path: str):
|
||||
"""Take an image with the webcam and save it to the specified
|
||||
path.
|
||||
|
||||
:param str img_path: path image should be saved to
|
||||
:returns: captured image
|
||||
|
||||
"""
|
||||
input = jetson.utils.videoSource('csi://0')
|
||||
output = jetson.utils.videoOutput(img_path)
|
||||
img = input.Capture()
|
||||
output.Render(img)
|
||||
return img
|
||||
|
||||
|
||||
|
||||
@app.route('/')
|
||||
def index():
|
||||
# TODO: call script and save initial image with bounding boxes
|
||||
# TODO: get predictions and output them in JSON via API
|
||||
# TODO: periodically get image from webcam and go to beginning
|
||||
# TODO: JSON format: [Nr, state (1-10), timestamp, time since below 3]
|
||||
return 'Server works'
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
scheduler.add_job(func=get_pred, trigger='interval', minutes=30)
|
||||
scheduler.start()
|
||||
atexit.register(scheduler.shutdown())
|
||||
app.run()
|
||||
Reference in New Issue
Block a user