potok/opencv
1
0
Fork 0
mirror of https://github.com/opencv/opencv.git synced 2026-07-31 08:13:04 +04:00
Code Packages Releases Wiki Activity

Merge pull request #25519 from gursimarsingh:improved_classification_sample

Browse Source 
Improved classification sample #25519

#25006 #25314

This pull requests replaces the caffe model for classification with onnx versions. It also adds resnet in model.yml. 

### Pull Request Readiness Checklist

See details at https://github.com/opencv/opencv/wiki/How_to_contribute#making-a-good-pull-request

- [x] I agree to contribute to the project under Apache 2 License.
- [x] To the best of my knowledge, the proposed patch is not based on a code under GPL or another license that is incompatible with OpenCV
- [x] The PR is proposed to the proper branch
- [x] There is a reference to the original bug report and related work
- [ ] There is accuracy test, performance test and test data in opencv_extra repository, if applicable
      Patch to opencv_extra has the same branch name.
- [x] The feature is well documented and sample code can be built with the project CMake
This commit is contained in:
Gursimar Singh
2024-08-06 11:46:11 +05:30
committed by GitHub
parent 3dcc8c38b4
commit 35eba9ca90
7 changed files with 493 additions and 240 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samples/dnn/classification.py
+115 -60
View File
@@ -1,49 +1,55 @@
import os
import glob
import argparse
import cv2 as cv
import numpy as np
import sys
from common import *
def help():
print(
'''
Firstly, download required models using `download_models.py` (if not already done). Set environment variable OPENCV_DOWNLOAD_CACHE_DIR to specify where models should be downloaded. Also, point OPENCV_SAMPLES_DATA_PATH to opencv/samples/data.\n"\n
To run:
python classification.py model_name --input=path/to/your/input/image/or/video (don't give --input flag if want to use device camera)
Sample command:
python classification.py googlenet --input=path/to/image
Model path can also be specified using --model argument
'''
)
def get_args_parser(func_args):
backends = (cv.dnn.DNN_BACKEND_DEFAULT, cv.dnn.DNN_BACKEND_INFERENCE_ENGINE,
cv.dnn.DNN_BACKEND_OPENCV, cv.dnn.DNN_BACKEND_VKCOM, cv.dnn.DNN_BACKEND_CUDA)
targets = (cv.dnn.DNN_TARGET_CPU, cv.dnn.DNN_TARGET_OPENCL, cv.dnn.DNN_TARGET_OPENCL_FP16, cv.dnn.DNN_TARGET_MYRIAD,
cv.dnn.DNN_TARGET_HDDL, cv.dnn.DNN_TARGET_VULKAN, cv.dnn.DNN_TARGET_CUDA, cv.dnn.DNN_TARGET_CUDA_FP16)
backends = ("default", "openvino", "opencv", "vkcom", "cuda")
targets = ("cpu", "opencl", "opencl_fp16", "ncs2_vpu", "hddl_vpu", "vulkan", "cuda", "cuda_fp16")
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument('--zoo', default=os.path.join(os.path.dirname(os.path.abspath(__file__)), 'models.yml'),
help='An optional path to file with preprocessing parameters.')
parser.add_argument('--input',
help='Path to input image or video file. Skip this argument to capture frames from a camera.')
parser.add_argument('--framework', choices=['caffe', 'tensorflow', 'darknet'],
help='Optional name of an origin framework of the model. '
'Detect it automatically if it does not set.')
parser.add_argument('--std', nargs='*', type=float,
help='Preprocess input image by dividing on a standard deviation.')
parser.add_argument('--crop', type=bool, default=False,
help='Preprocess input image by dividing on a standard deviation.')
parser.add_argument('--initial_width', type=int,
help='Preprocess input image by initial resizing to a specific width.')
parser.add_argument('--initial_height', type=int,
help='Preprocess input image by initial resizing to a specific height.')
parser.add_argument('--backend', choices=backends, default=cv.dnn.DNN_BACKEND_DEFAULT, type=int,
help="Choose one of computation backends: "
"%d: automatically (by default), "
"%d: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
"%d: OpenCV implementation, "
"%d: VKCOM, "
"%d: CUDA" % backends)
parser.add_argument('--target', choices=targets, default=cv.dnn.DNN_TARGET_CPU, type=int,
help='Choose one of target computation devices: '
'%d: CPU target (by default), '
'%d: OpenCL, '
'%d: OpenCL fp16 (half-float precision), '
'%d: NCS2 VPU, '
'%d: HDDL VPU, '
'%d: Vulkan, '
'%d: CUDA, '
'%d: CUDA fp16 (half-float preprocess)'% targets)
help='Center crop the image.')
parser.add_argument('--backend', default="default", type=str, choices=backends,
help="Choose one of computation backends: "
"default: automatically (by default), "
"openvino: Intel's Deep Learning Inference Engine (https://software.intel.com/openvino-toolkit), "
"opencv: OpenCV implementation, "
"vkcom: VKCOM, "
"cuda: CUDA, "
"webnn: WebNN")
parser.add_argument('--target', default="cpu", type=str, choices=targets,
help="Choose one of target computation devices: "
"cpu: CPU target (by default), "
"opencl: OpenCL, "
"opencl_fp16: OpenCL fp16 (half-float precision), "
"ncs2_vpu: NCS2 VPU, "
"hddl_vpu: HDDL VPU, "
"vulkan: Vulkan, "
"cuda: CUDA, "
"cuda_fp16: CUDA fp16 (half-float preprocess)")
args, _ = parser.parse_known_args()
add_preproc_args(args.zoo, parser, 'classification')
@@ -52,41 +58,76 @@ def get_args_parser(func_args):
formatter_class=argparse.ArgumentDefaultsHelpFormatter)
return parser.parse_args(func_args)
def load_images(directory):
# List all common image file extensions, feel free to add more if needed
extensions = ['jpg', 'jpeg', 'png', 'bmp', 'tif', 'tiff']
files = []
for extension in extensions:
files.extend(glob.glob(os.path.join(directory, f'*.{extension}')))
return files
def main(func_args=None):
args = get_args_parser(func_args)
args.model = findFile(args.model)
args.config = findFile(args.config)
args.classes = findFile(args.classes)
if args.alias is None or hasattr(args, 'help'):
help()
exit(1)
args.model = findModel(args.model, args.sha1)
args.labels = findFile(args.labels)
# Load names of classes
classes = None
if args.classes:
with open(args.classes, 'rt') as f:
classes = f.read().rstrip('\n').split('\n')
labels = None
if args.labels:
with open(args.labels, 'rt') as f:
labels = f.read().rstrip('\n').split('\n')
# Load a network
net = cv.dnn.readNet(args.model, args.config, args.framework)
net.setPreferableBackend(args.backend)
net.setPreferableTarget(args.target)
net = cv.dnn.readNet(args.model)
net.setPreferableBackend(get_backend_id(args.backend))
net.setPreferableTarget(get_target_id(args.target))
winName = 'Deep learning image classification in OpenCV'
cv.namedWindow(winName, cv.WINDOW_NORMAL)
cap = cv.VideoCapture(args.input if args.input else 0)
isdir = False
if args.input:
input_path = args.input
if os.path.isdir(input_path):
isdir = True
image_files = load_images(input_path)
if not image_files:
print("No images found in the directory.")
exit(-1)
current_image_index = 0
else:
input_path = findFile(input_path)
cap = cv.VideoCapture(input_path)
if not cap.isOpened():
print("Failed to open the input video")
exit(-1)
else:
cap = cv.VideoCapture(0)
while cv.waitKey(1) < 0:
hasFrame, frame = cap.read()
if not hasFrame:
cv.waitKey()
break
if isdir:
if current_image_index >= len(image_files):
break
frame = cv.imread(image_files[current_image_index])
current_image_index += 1
else:
hasFrame, frame = cap.read()
if not hasFrame:
cv.waitKey()
break
# Create a 4D blob from a frame.
inpWidth = args.width if args.width else frame.shape[1]
inpHeight = args.height if args.height else frame.shape[0]
if args.initial_width and args.initial_height:
frame = cv.resize(frame, (args.initial_width, args.initial_height))
blob = cv.dnn.blobFromImage(frame, args.scale, (inpWidth, inpHeight), args.mean, args.rgb, crop=args.crop)
if args.std:
blob[0] /= np.asarray(args.std, dtype=np.float32).reshape(3, 1, 1)
@@ -95,22 +136,36 @@ def main(func_args=None):
net.setInput(blob)
out = net.forward()
# Get a class with a highest score.
out = out.flatten()
classId = np.argmax(out)
confidence = out[classId]
(h, w, _) = frame.shape
roi_rows = min(300, h)
roi_cols = min(1000, w)
frame[:roi_rows,:roi_cols,:] >>= 1
# Put efficiency information.
t, _ = net.getPerfProfile()
label = 'Inference time: %.2f ms' % (t * 1000.0 / cv.getTickFrequency())
cv.putText(frame, label, (0, 15), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0))
label = 'Inference time: %.1f ms' % (t * 1000.0 / cv.getTickFrequency())
cv.putText(frame, label, (15, 30), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0))
# Print predicted class.
label = '%s: %.4f' % (classes[classId] if classes else 'Class #%d' % classId, confidence)
cv.putText(frame, label, (0, 40), cv.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0))
# Print predicted classes.
out = out.flatten()
K = 5
topKidx = np.argpartition(out, -K)[-K:]
for i in range(K):
classId = topKidx[i]
confidence = out[classId]
label = '%s: %.2f' % (labels[classId] if labels else 'Class #%d' % classId, confidence)
cv.putText(frame, label, (15, 90 + i*30), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0))
cv.imshow(winName, frame)
key = cv.waitKey(1000 if isdir else 100)
if key >= 0:
key &= 255
if key == ord(' '):
key = cv.waitKey() & 255
if key == ord('q') or key == 27: # Wait for 1 second on each image, press 'q' to exit
sys.exit(0)
cv.waitKey()
if __name__ == "__main__":
main()
main()