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 os
import glob import glob
import random import random
import darknet
import time import time
import cv2 import cv2
import numpy as np import numpy as np
@ -97,14 +96,17 @@ def prepare_batch(images, network, channels=3):
return darknet.IMAGE(width, height, channels, darknet_images) 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. # Darknet doesn't accept numpy images.
# Create one with image we reuse for each detect # Create one with image we reuse for each detect
width = darknet.network_width(network) width = darknet.network_width(network)
height = darknet.network_height(network) height = darknet.network_height(network)
darknet_image = darknet.make_image(width, height, 3) 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_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image_resized = cv2.resize(image_rgb, (width, height), image_resized = cv2.resize(image_rgb, (width, height),
interpolation=cv2.INTER_LINEAR) interpolation=cv2.INTER_LINEAR)