Skip to content

Auto download model and some fixes #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
mmoollllee wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Auto download model and some fixes #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
27ca3f7
Download Model if not exists
mmoollllee Aug 8, 2023
f387f39
Fix CV2 version check for newer than 4.6.0
mmoollllee Aug 8, 2023
a95cba9
Add Zenity Requirement
mmoollllee Aug 8, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,8 @@ It uses YOLOv3-608 weights to pre-annotate the bounding boxes in images. It redu
* '/output' : Sample of output files after annotation.

## Usage:
requires zenity (dialog boxes for the command-line) to work

```
python yolo_annotation_tool.py
```
Expand Down
18 changes: 9 additions & 9 deletions recognize_objects.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,10 @@ def __init__(self, yolo_model_weights, yolo_model_cfg, coco_names):
self.net = cv2.dnn.readNetFromDarknet(self.cfg, self.weights)
self.ln = self.net.getLayerNames()

if cv2.__version__ == '4.6.0':
self.ln = [self.ln[i - 1] for i in self.net.getUnconnectedOutLayers()]
else:
self.ln = [self.ln[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]
if cv2.__version__ >= '4.6.0':
self.ln = [self.ln[i - 1] for i in self.net.getUnconnectedOutLayers()]
else:
self.ln = [self.ln[i[0] - 1] for i in self.net.getUnconnectedOutLayers()]

def process_frame(self, frame, all_labels, show = True):

Expand Down Expand Up @@ -91,11 +91,11 @@ def process_frame(self, frame, all_labels, show = True):
if len(idxs) > 0:
print("Objects detected in this frame:\n")
for i in idxs.flatten():
if cv2.__version__ == '4.6.0':
i = i
else:
i = i[0]
print(self.classes[classIDs[i]])
if cv2.__version__ >= '4.6.0':
i = i
else:
i = i[0]
print(self.classes[classIDs[i]])
print("\n###################################\n")

# If it is person, cat, dog detected
Expand Down
11 changes: 10 additions & 1 deletion yolo_annotation_tool.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import threading
import csv
from recognize_objects import recognize_objects
from urllib.request import urlretrieve


import configparser
Expand All @@ -20,7 +21,15 @@

all_labels = []

ob = recognize_objects("./models/yolov3.weights", "./models/yolov3.cfg","./models/coco.names")
weights_file = "./models/yolov3.weights"
if not os.path.isfile(weights_file):
print("Downloading weights file...")
urlretrieve("https://pjreddie.com/media/files/yolov3.weights", weights_file)

ob = recognize_objects(weights_file, "./models/yolov3.cfg","./models/coco.names")


#if

colors = [(255, 0, 0), (255, 255, 0), (0, 255, 0), (0, 255, 255), (0, 0, 255), (255, 0, 255),
(255, 100, 0), (255, 255, 100), (0, 255, 100), (50, 255, 255), (50, 0, 255), (255, 50, 255),
Expand Down