Skip to content

Instantly share code, notes, and snippets.

@pythonlessons

pythonlessons/Collect_training_data.py

Last active September 29, 2020 08:29
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save pythonlessons/201ddd98f01b73a4682cd20e023245b5 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save pythonlessons/201ddd98f01b73a4682cd20e023245b5 to your computer and use it in GitHub Desktop.
Download ZIP
Collect_training_data
Raw
Collect_training_data.py
offset = 30
sct = mss.mss()
yolo = Load_Yolo_model()
while True:
img = np.array(sct.grab({"top": 87-offset, "left": 1920, "width": 1280, "height": 720, "mon": -1}))
img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
image, bboxes = detect_enemy(yolo, np.copy(img), input_size=YOLO_INPUT_SIZE, CLASSES=TRAIN_CLASSES, rectangle_colors=(255,0,0))
if len(bboxes) > 0:
CreateXMLfile("XML_Detections", str(int(time.time())), img, bboxes, read_class_names(TRAIN_CLASSES))
print("got it")
time.sleep(2)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment