Allow image_detection to receive a cv2 image; remove duplicate darknet import (#8387)

This commit is contained in:
Scott Martin 2022-05-16 19:04:54 -04:00 committed by GitHub
parent 695d562918
commit 4a7c2802e0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 5 additions and 3 deletions

View File

@ -2,7 +2,6 @@ import argparse
import os
import glob
import random
import darknet
import time
import cv2
import numpy as np
@ -97,14 +96,17 @@ def prepare_batch(images, network, channels=3):
return darknet.IMAGE(width, height, channels, darknet_images)
def image_detection(image_path, network, class_names, class_colors, thresh):
def image_detection(image_or_path, network, class_names, class_colors, thresh):
# Darknet doesn't accept numpy images.
# Create one with image we reuse for each detect
width = darknet.network_width(network)
height = darknet.network_height(network)
darknet_image = darknet.make_image(width, height, 3)
image = cv2.imread(image_path)
if type(image_or_path) == "str":
image = cv2.imread(image_or_path)
else:
image = image_or_path
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_resized = cv2.resize(image_rgb, (width, height),
interpolation=cv2.INTER_LINEAR)