zenon/master-thesis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add tkinter import for commandline opencv

Browse Source
This commit is contained in:
Tobias Eidelpes 2023-02-08 10:05:09 +01:00
parent dfccc908c3
commit 2e0d8b9364
1 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
jetson-deployment/model.py
View File

@ -2,6 +2,7 @@ import argparse
import cv2 import cv2
import torch import torch
from torchvision import transforms from torchvision import transforms
import tkinter
def load_models(yolo_path: str, resnet_path: str): def load_models(yolo_path: str, resnet_path: str):
@ -39,6 +40,7 @@ def detect(img_path: str, yolo_path: str, resnet_path: str):
box_coords = get_boxes(first_stage, img) box_coords = get_boxes(first_stage, img)
box_coords.sort_values(by=['xmin'], ignore_index=True, inplace=True) box_coords.sort_values(by=['xmin'], ignore_index=True, inplace=True)
print(box_coords)
predictions = {} predictions = {}
for idx, row in box_coords.iterrows(): for idx, row in box_coords.iterrows():
@ -61,6 +63,9 @@ def detect(img_path: str, yolo_path: str, resnet_path: str):
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 255, 255),
2, cv2.LINE_AA) 2, cv2.LINE_AA)
cv2.imshow('original', original)
cv2.waitKey(0)
cv2.destroyAllWindows()
return (original, predictions) return (original, predictions)
@ -161,3 +166,5 @@ if __name__ == '__main__':
detect(opt.source, 'yolo.pt', 'resnet.pt') detect(opt.source, 'yolo.pt', 'resnet.pt')
if opt.onnx: if opt.onnx:
export_to_onnx('yolo.pt', 'resnet.pt') export_to_onnx('yolo.pt', 'resnet.pt')