Commit cf02452d authored by Tom Pillot's avatar Tom Pillot
Browse files

Refactor some repeated code

parent d68d6c01
import apriltag
import cv2
from game.webcam_chooser import WebcamWindow
class Detector(apriltag.Detector):
DECISION_MARGIN = 40
def __init__(self, cam):
def __init__(self):
super().__init__()
self.cam = cam
def detect(self, img, return_image=False):
dets = super().detect(img)
......@@ -20,11 +21,11 @@ class Detector(apriltag.Detector):
# Refresh cam before trying to detect tags, else it reads an old image
for i in range(5):
self.cam.read()
WebcamWindow.cam.read()
tag_id = None
while not tag_id:
ret, frame = self.cam.read()
ret, frame = WebcamWindow.cam.read()
if not ret:
print("Failed to grab frame.")
break
......@@ -55,6 +56,3 @@ class Detector(apriltag.Detector):
pos = det.center.astype(int) + (-10, 10)
cv2.putText(img, ident, tuple(pos), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
return dets
def stop_webcam(self):
self.cam.release()
import cv2
from socket import *
from PyQt5 import QtWidgets as qtw
from PyQt5 import QtCore as qtc
from PyQt5 import QtGui as qtg
from game.scoreboard_ui import ScoreWindowUi
from PyQt5 import QtWidgets as qtw
from game.detector import Detector
from game.scoreboard_ui import ScoreWindowUi
from game.webcam_chooser import WebcamWindow, frame_to_pixmap
class ScoreWindow(ScoreWindowUi):
......@@ -13,14 +15,14 @@ class ScoreWindow(ScoreWindowUi):
PORT = 42563
MAX_START_SENT = 5
def __init__(self, teams, cam):
def __init__(self, teams):
super().__init__()
self.teams = teams
self.initial_time = 0
self.time_is_red = False
self.join_started = False
self.detector = Detector(cam)
self.detector = Detector()
self.socket = socket(AF_INET, SOCK_DGRAM)
self.socket.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
......@@ -94,16 +96,10 @@ class ScoreWindow(ScoreWindowUi):
def refresh_webcam(self, detect_tags=True):
# Read an image from the webcam and detect the tags it contains
ret, frame = self.detector.cam.read()
ret, frame = WebcamWindow.cam.read()
dets = self.detector.show_tags_on_image(frame)
if ret:
frame = cv2.resize(frame, (640, 480))
# Convert frame to QImage
cv2.cvtColor(frame, cv2.COLOR_RGB2BGR, frame)
rows, cols, _ = frame.shape
qimage = qtg.QImage(frame.data, cols, rows, qtg.QImage.Format_RGB888)
pixmap = qtg.QPixmap()
pixmap.convertFromImage(qimage, qtc.Qt.AutoColor)
pixmap = frame_to_pixmap(frame)
self.webcam_label.setPixmap(pixmap)
else:
print("Failed to grab frame.")
......
from PyQt5 import QtWidgets as qtw
from PyQt5 import QtCore as qtc
from PyQt5 import QtGui as qtg
from PyQt5 import QtWidgets as qtw
class ScoreWindowUi(qtw.QMainWindow):
......
from PyQt5 import QtCore as qtc
from PyQt5 import QtGui as qtg
from game.detector import Detector
from game.team_chooser_ui import TeamWindowUi
from game.team import Team
from game.team_chooser_ui import TeamWindowUi
class TeamWindow(TeamWindowUi):
......@@ -16,10 +17,10 @@ class TeamWindow(TeamWindowUi):
96: "#FE005D", 97: "#FE002A", 98: "#474747", 99: "#8C8C8C", 100: "#643030", 101: "#9A6565", 102: "#BC9F9F"
}
def __init__(self, cam):
def __init__(self):
super().__init__()
self.detector = Detector(cam)
self.detector = Detector()
self.teams = {} # The keys are the tag ids and the values are the teams
self.add_activities()
......
from PyQt5 import QtWidgets as qtw
from PyQt5 import QtCore as qtc
from PyQt5 import QtGui as qtg
from PyQt5 import QtWidgets as qtw
class AlignDelegate(qtw.QStyledItemDelegate):
......
import cv2
from PyQt5 import QtWidgets as qtw
from PyQt5 import QtCore as qtc
from PyQt5 import QtGui as qtg
from PyQt5 import QtWidgets as qtw
class WebcamWindow(qtw.QWidget):
WEBCAM_REFRESH_TIME = 50
WEBCAM_SIZE = (640, 480)
cam = None
def __init__(self):
super().__init__()
......@@ -31,6 +33,7 @@ class WebcamWindow(qtw.QWidget):
self.camera_indexes = return_camera_indexes()
self.current_index = 0
self.cam = cv2.VideoCapture(self.camera_indexes[0])
self.update_cam(self.cam)
self.left_button.clicked.connect(self.show_previous)
self.right_button.clicked.connect(self.show_next)
......@@ -42,13 +45,7 @@ class WebcamWindow(qtw.QWidget):
# Read an image from the webcam and detect the tags it contains
ret, frame = self.cam.read()
if ret:
frame = cv2.resize(frame, (640, 480))
# Convert frame to QImage
cv2.cvtColor(frame, cv2.COLOR_RGB2BGR, frame)
rows, cols, _ = frame.shape
qimage = qtg.QImage(frame.data, cols, rows, qtg.QImage.Format_RGB888)
pixmap = qtg.QPixmap()
pixmap.convertFromImage(qimage, qtc.Qt.AutoColor)
pixmap = frame_to_pixmap(frame)
self.webcam_label.setPixmap(pixmap)
else:
print("Failed to grab frame.")
......@@ -68,8 +65,11 @@ class WebcamWindow(qtw.QWidget):
self.cam.release()
new_index = self.camera_indexes[self.current_index]
self.cam = cv2.VideoCapture(new_index)
self.cam.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
self. cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
self.update_cam(self.cam)
@classmethod
def update_cam(cls, new_cam):
cls.cam = new_cam
def return_camera_indexes():
......@@ -85,3 +85,14 @@ def return_camera_indexes():
index += 1
i -= 1
return arr
def frame_to_pixmap(frame):
frame = cv2.resize(frame, WebcamWindow.WEBCAM_SIZE)
# Convert frame to QImage
cv2.cvtColor(frame, cv2.COLOR_RGB2BGR, frame)
rows, cols, _ = frame.shape
qimage = qtg.QImage(frame.data, cols, rows, qtg.QImage.Format_RGB888)
pixmap = qtg.QPixmap()
pixmap.convertFromImage(qimage, qtc.Qt.AutoColor)
return pixmap
import sys
from PyQt5 import QtWidgets as qtw
from game.team_chooser import TeamWindow
from game.scoreboard import ScoreWindow
from game.team_chooser import TeamWindow
from game.webcam_chooser import WebcamWindow
......@@ -11,28 +13,17 @@ class MainWindow(qtw.QWidget):
self.stacked_layout = qtw.QStackedLayout()
self.webcam_window = WebcamWindow()
self.team_window = TeamWindow()
self.scoreboard_window = ScoreWindow(self.team_window.teams)
self.stacked_layout.addWidget(self.webcam_window)
self.stacked_layout.addWidget(self.team_window)
self.stacked_layout.addWidget(self.scoreboard_window)
# We need to wait for the webcam to be chosen before creating the other windows
self.webcam_window.start_button.clicked.connect(self.show_team_window)
def show_team_window(self):
self.team_window = TeamWindow(self.webcam_window.cam)
self.scoreboard_window = ScoreWindow(self.team_window.teams, self.webcam_window.cam)
# Add teams to test the application
if len(sys.argv) > 1 and sys.argv[1] == "test":
for tag_id in w.team_window.team_colors.keys():
tag_str = str(tag_id)
name = "Test " + str(tag_id)
w.team_window.name_input.setText(name)
w.team_window.tag_label.setText(tag_str)
w.team_window.add_button.click()
self.team_window.start_button.clicked.connect(self.start_activity)
self.stacked_layout.addWidget(self.team_window)
self.stacked_layout.addWidget(self.scoreboard_window)
def show_team_window(self):
self.stacked_layout.setCurrentIndex(1)
def start_activity(self):
......@@ -47,4 +38,14 @@ class MainWindow(qtw.QWidget):
if __name__ == '__main__':
app = qtw.QApplication(sys.argv)
w = MainWindow()
# Add teams to test the application
if len(sys.argv) > 1 and sys.argv[1] == "test":
for tag_id in w.team_window.team_colors.keys():
tag_str = str(tag_id)
name = "Test " + str(tag_id)
w.team_window.name_input.setText(name)
w.team_window.tag_label.setText(tag_str)
w.team_window.add_button.click()
sys.exit(app.exec_())
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment