Skip to content

Commit 2ade3ba

Browse files
digas99digas99
committed
created class Detector
1 parent 8ffd2de commit 2ade3ba

File tree

9 files changed

+134
-117
lines changed

9 files changed

+134
-117
lines changed

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,5 +5,6 @@ venv
55
*.weights
66
*.cfg
77
*.names
8+
yolov3
89

910
output

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ chmod +x get_yolov3.sh
2525
4. Run the app
2626

2727
```bash
28-
python3 app.py -i <image> -c yolov3.cfg -w yolov3.weights -cl coco.names
28+
python3 src/app.py -i <image> -c yolov3/yolov3.cfg -w yolov3/yolov3.weights -cl yolov3/coco.names
2929
```
3030
or
3131
```bash

‎app.py‎

Lines changed: 0 additions & 115 deletions
This file was deleted.

‎get_yolov3.sh‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,7 @@
1+
#!/bin/bash
2+
mkdir -p yolov3
3+
cd yolov3
4+
15
wget https://pjreddie.com/media/files/yolov3.weights
26
wget https://raw.githubusercontent.com/pjreddie/darknet/master/cfg/yolov3.cfg
37
wget https://raw.githubusercontent.com/pjreddie/darknet/master/data/coco.names

‎run.sh‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1,5 @@
1-
python3 app.py -i $1 -c yolov3.cfg -w yolov3.weights -cl coco.names
1+
#!/bin/bash
2+
3+
SOURCE=src
4+
DATA=yolov3
5+
python3 $SOURCE/app.py -i $1 -c $DATA/yolov3.cfg -w $DATA/yolov3.weights -cl $DATA/coco.names

‎src/__pycache__/detection.cpython-310.pyc‎

2.57 KB
Binary file not shown.

‎src/app.py‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
import cv2
2+
import argparse
3+
import numpy as np
4+
import os
5+
6+
from detection import Detector
7+
8+
classes = None
9+
COLORS = None
10+
11+
BASE_DIR = os.path.dirname(os.path.abspath(__name__))
12+
13+
def main():
14+
# handle command line arguments
15+
ap = argparse.ArgumentParser()
16+
ap.add_argument('-i', '--image', required=True,
17+
help = 'path to input image')
18+
ap.add_argument('-c', '--config', required=True,
19+
help = 'path to yolo config file')
20+
ap.add_argument('-w', '--weights', required=True,
21+
help = 'path to yolo pre-trained weights')
22+
ap.add_argument('-cl', '--classes', required=True,
23+
help = 'path to text file containing class names')
24+
args = ap.parse_args()
25+
26+
detector = Detector(args.config, args.weights, args.classes)
27+
28+
image = cv2.imread(args.image)
29+
image = detector.detect(image)
30+
31+
output_dir = os.path.join(BASE_DIR, "output")
32+
if not os.path.exists(output_dir):
33+
os.makedirs(output_dir)
34+
35+
num_files = len(os.listdir(output_dir))
36+
detector.export(image, f"{output_dir}/output{num_files}.jpg")
37+
38+
detector.show(image)
39+
40+
41+
if __name__ == "__main__":
42+
main()

‎src/camera.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
import cv2
2+
3+
video = cv2.VideoCapture(0)
4+
5+
while True:
6+
ret, image = video.read()
7+
cv2.imshow('image', image)
8+
if cv2.waitKey(1) & 0xFF == ord('q'):
9+
break
10+
11+
video.release()
12+
13+
cv2.destroyAllWindows()

‎src/detection.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
import cv2
2+
import numpy as np
3+
4+
class Detector:
5+
def __init__(self, config, weights, classes):
6+
self.net = cv2.dnn.readNet(weights, config)
7+
self.output_layers = self.net.getUnconnectedOutLayersNames()
8+
self.scale = 0.00392
9+
self.conf_threshold = 0.5
10+
self.nms_threshold = 0.4
11+
12+
with open(classes, 'r') as f:
13+
self.classes = [line.strip() for line in f.readlines()]
14+
15+
self.colors = np.random.uniform(0, 255, size=(len(self.classes), 3))
16+
17+
def detect(self, image):
18+
H, W, _ = image.shape
19+
blob = cv2.dnn.blobFromImage(image, self.scale, (416,416), (0,0,0), True, crop=False)
20+
self.net.setInput(blob)
21+
outs = self.net.forward(self.output_layers)
22+
23+
class_ids = []
24+
confidences = []
25+
boxes = []
26+
27+
for out in outs:
28+
for detection in out:
29+
scores = detection[5:]
30+
class_id = np.argmax(scores)
31+
confidence = scores[class_id]
32+
if confidence > self.conf_threshold:
33+
center_x = int(detection[0] * W)
34+
center_y = int(detection[1] * H)
35+
w = int(detection[2] * W)
36+
h = int(detection[3] * H)
37+
x = center_x - w / 2
38+
y = center_y - h / 2
39+
class_ids.append(class_id)
40+
confidences.append(float(confidence))
41+
boxes.append([x, y, w, h])
42+
43+
indices = cv2.dnn.NMSBoxes(boxes, confidences, self.conf_threshold, self.nms_threshold)
44+
45+
for i in indices:
46+
box = boxes[i]
47+
x = box[0]
48+
y = box[1]
49+
w = box[2]
50+
h = box[3]
51+
52+
self.draw_bounding_box(image, class_ids[i], round(x), round(y), round(x+w), round(y+h))
53+
54+
return image
55+
56+
def draw_bounding_box(self, img, class_id, x, y, x_plus_w, y_plus_h):
57+
label = str(self.classes[class_id])
58+
color = self.colors[class_id]
59+
cv2.rectangle(img, (x,y), (x_plus_w,y_plus_h), color, 2)
60+
cv2.putText(img, label, (x-10,y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
61+
62+
def show(self, image):
63+
cv2.imshow('Object Detection', image)
64+
cv2.waitKey(0)
65+
cv2.destroyAllWindows()
66+
67+
def export(self, image, path):
68+
cv2.imwrite(path, image)

0 commit comments

Comments
 (0)