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

Refactor some repeated code

parent d68d6c01
import apriltag import apriltag
import cv2 import cv2
from game.webcam_chooser import WebcamWindow
class Detector(apriltag.Detector): class Detector(apriltag.Detector):
DECISION_MARGIN = 40 DECISION_MARGIN = 40
def __init__(self, cam): def __init__(self):
super().__init__() super().__init__()
self.cam = cam
def detect(self, img, return_image=False): def detect(self, img, return_image=False):
dets = super().detect(img) dets = super().detect(img)
...@@ -20,11 +21,11 @@ class Detector(apriltag.Detector): ...@@ -20,11 +21,11 @@ class Detector(apriltag.Detector):
# Refresh cam before trying to detect tags, else it reads an old image # Refresh cam before trying to detect tags, else it reads an old image
for i in range(5): for i in range(5):
self.cam.read() WebcamWindow.cam.read()
tag_id = None tag_id = None
while not tag_id: while not tag_id:
ret, frame = self.cam.read() ret, frame = WebcamWindow.cam.read()
if not ret: if not ret:
print("Failed to grab frame.") print("Failed to grab frame.")
break break
...@@ -55,6 +56,3 @@ class Detector(apriltag.Detector): ...@@ -55,6 +56,3 @@ class Detector(apriltag.Detector):
pos = det.center.astype(int) + (-10, 10) pos = det.center.astype(int) + (-10, 10)
cv2.putText(img, ident, tuple(pos), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2) cv2.putText(img, ident, tuple(pos), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
return dets return dets
def stop_webcam(self):
self.cam.release()
import cv2
from socket import * from socket import *
from PyQt5 import QtWidgets as qtw
from PyQt5 import QtCore as qtc from PyQt5 import QtCore as qtc
from PyQt5 import QtGui as qtg 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.detector import Detector
from game.scoreboard_ui import ScoreWindowUi
from game.webcam_chooser import WebcamWindow, frame_to_pixmap
class ScoreWindow(ScoreWindowUi): class ScoreWindow(ScoreWindowUi):
...@@ -13,14 +15,14 @@ class ScoreWindow(ScoreWindowUi): ...@@ -13,14 +15,14 @@ class ScoreWindow(ScoreWindowUi):
PORT = 42563 PORT = 42563
MAX_START_SENT = 5 MAX_START_SENT = 5
def __init__(self, teams, cam): def __init__(self, teams):
super().__init__() super().__init__()
self.teams = teams self.teams = teams
self.initial_time = 0 self.initial_time = 0
self.time_is_red = False self.time_is_red = False
self.join_started = False self.join_started = False
self.detector = Detector(cam) self.detector = Detector()
self.socket = socket(AF_INET, SOCK_DGRAM) self.socket = socket(AF_INET, SOCK_DGRAM)
self.socket.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1) self.socket.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
...@@ -94,16 +96,10 @@ class ScoreWindow(ScoreWindowUi): ...@@ -94,16 +96,10 @@ class ScoreWindow(ScoreWindowUi):
def refresh_webcam(self, detect_tags=True): def refresh_webcam(self, detect_tags=True):
# Read an image from the webcam and detect the tags it contains # 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) dets = self.detector.show_tags_on_image(frame)
if ret: if ret:
frame = cv2.resize(frame, (640, 480)) pixmap = frame_to_pixmap(frame)
# 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)
self.webcam_label.setPixmap(pixmap) self.webcam_label.setPixmap(pixmap)
else: else:
print("Failed to grab frame.") print("Failed to grab frame.")
......
from PyQt5 import QtWidgets as qtw
from PyQt5 import QtCore as qtc from PyQt5 import QtCore as qtc
from PyQt5 import QtGui as qtg from PyQt5 import QtGui as qtg
from PyQt5 import QtWidgets as qtw
class ScoreWindowUi(qtw.QMainWindow): class ScoreWindowUi(qtw.QMainWindow):
......
from PyQt5 import QtCore as qtc from PyQt5 import QtCore as qtc
from PyQt5 import QtGui as qtg from PyQt5 import QtGui as qtg
from game.detector import Detector from game.detector import Detector
from game.team_chooser_ui import TeamWindowUi
from game.team import Team from game.team import Team
from game.team_chooser_ui import TeamWindowUi
class TeamWindow(TeamWindowUi): class TeamWindow(TeamWindowUi):
...@@ -16,10 +17,10 @@ class TeamWindow(TeamWindowUi): ...@@ -16,10 +17,10 @@ class TeamWindow(TeamWindowUi):
96: "#FE005D", 97: "#FE002A", 98: "#474747", 99: "#8C8C8C", 100: "#643030", 101: "#9A6565", 102: "#BC9F9F" 96: "#FE005D", 97: "#FE002A", 98: "#474747", 99: "#8C8C8C", 100: "#643030", 101: "#9A6565", 102: "#BC9F9F"
} }
def __init__(self, cam): def __init__(self):
super().__init__() super().__init__()
self.detector = Detector(cam) self.detector = Detector()
self.teams = {} # The keys are the tag ids and the values are the teams self.teams = {} # The keys are the tag ids and the values are the teams
self.add_activities() self.add_activities()
......
from PyQt5 import QtWidgets as qtw
from PyQt5 import QtCore as qtc from PyQt5 import QtCore as qtc
from PyQt5 import QtGui as qtg from PyQt5 import QtGui as qtg
from PyQt5 import QtWidgets as qtw
class AlignDelegate(qtw.QStyledItemDelegate): class AlignDelegate(qtw.QStyledItemDelegate):
......
import cv2 import cv2
from PyQt5 import QtWidgets as qtw
from PyQt5 import QtCore as qtc from PyQt5 import QtCore as qtc
from PyQt5 import QtGui as qtg from PyQt5 import QtGui as qtg
from PyQt5 import QtWidgets as qtw
class WebcamWindow(qtw.QWidget): class WebcamWindow(qtw.QWidget):
WEBCAM_REFRESH_TIME = 50 WEBCAM_REFRESH_TIME = 50
WEBCAM_SIZE = (640, 480)
cam = None
def __init__(self): def __init__(self):
super().__init__() super().__init__()
...@@ -31,6 +33,7 @@ class WebcamWindow(qtw.QWidget): ...@@ -31,6 +33,7 @@ class WebcamWindow(qtw.QWidget):
self.camera_indexes = return_camera_indexes() self.camera_indexes = return_camera_indexes()
self.current_index = 0 self.current_index = 0
self.cam = cv2.VideoCapture(self.camera_indexes[0]) self.cam = cv2.VideoCapture(self.camera_indexes[0])
self.update_cam(self.cam)
self.left_button.clicked.connect(self.show_previous) self.left_button.clicked.connect(self.show_previous)
self.right_button.clicked.connect(self.show_next) self.right_button.clicked.connect(self.show_next)
...@@ -42,13 +45,7 @@ class WebcamWindow(qtw.QWidget): ...@@ -42,13 +45,7 @@ class WebcamWindow(qtw.QWidget):
# Read an image from the webcam and detect the tags it contains # Read an image from the webcam and detect the tags it contains
ret, frame = self.cam.read() ret, frame = self.cam.read()
if ret: if ret:
frame = cv2.resize(frame, (640, 480)) pixmap = frame_to_pixmap(frame)
# 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)
self.webcam_label.setPixmap(pixmap) self.webcam_label.setPixmap(pixmap)
else: else:
print("Failed to grab frame.") print("Failed to grab frame.")
...@@ -68,8 +65,11 @@ class WebcamWindow(qtw.QWidget): ...@@ -68,8 +65,11 @@ class WebcamWindow(qtw.QWidget):
self.cam.release() self.cam.release()
new_index = self.camera_indexes[self.current_index] new_index = self.camera_indexes[self.current_index]
self.cam = cv2.VideoCapture(new_index) self.cam = cv2.VideoCapture(new_index)
self.cam.set(cv2.CAP_PROP_FRAME_WIDTH, 640) self.update_cam(self.cam)
self. cam.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
@classmethod
def update_cam(cls, new_cam):
cls.cam = new_cam
def return_camera_indexes(): def return_camera_indexes():
...@@ -85,3 +85,14 @@ def return_camera_indexes(): ...@@ -85,3 +85,14 @@ def return_camera_indexes():
index += 1 index += 1
i -= 1 i -= 1
return arr 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 import sys
from PyQt5 import QtWidgets as qtw from PyQt5 import QtWidgets as qtw
from game.team_chooser import TeamWindow
from game.scoreboard import ScoreWindow from game.scoreboard import ScoreWindow
from game.team_chooser import TeamWindow
from game.webcam_chooser import WebcamWindow from game.webcam_chooser import WebcamWindow
...@@ -11,28 +13,17 @@ class MainWindow(qtw.QWidget): ...@@ -11,28 +13,17 @@ class MainWindow(qtw.QWidget):
self.stacked_layout = qtw.QStackedLayout() self.stacked_layout = qtw.QStackedLayout()
self.webcam_window = WebcamWindow() 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.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) 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.team_window.start_button.clicked.connect(self.start_activity)
self.stacked_layout.addWidget(self.team_window) def show_team_window(self):
self.stacked_layout.addWidget(self.scoreboard_window)
self.stacked_layout.setCurrentIndex(1) self.stacked_layout.setCurrentIndex(1)
def start_activity(self): def start_activity(self):
...@@ -47,4 +38,14 @@ class MainWindow(qtw.QWidget): ...@@ -47,4 +38,14 @@ class MainWindow(qtw.QWidget):
if __name__ == '__main__': if __name__ == '__main__':
app = qtw.QApplication(sys.argv) app = qtw.QApplication(sys.argv)
w = MainWindow() 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_()) sys.exit(app.exec_())
Supports Markdown
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