DylAFC.py

#!/usr/bin/python3.6
import os
import socket
import time
import numpy as np
from sys import argv
from tkinter import Tk, Frame, Label, E, W, Event, StringVar, Radiobutton, mainloop
from PIL import ImageTk, Image


class AFC:

	"""A Class for doing AFC studies. This one only does 2AFC. It is used for integrating with a tkinter interface as seen in DylAFC."""
	def __init__(self, posDir: str, negDir: str, ansDir: str, ip: str, port: str, n0: int, n1: int, logFile: str):
		"""Creates a new AFC class with the given parameters. Opens the file name provided for logging. Does not connect yet."""
		self.decision: int = -1
		self.ready: bool = True
		self.img1: int = -1
		self.img2: int = -1
		self.mode: str = 'training'
		self.posDir: str = posDir
		self.negDir: str = negDir
		self.ansDir: str = ansDir
		self.counter: int = 0
		self.imgIndex: int = 0
		self.ip: str = ip
		self.port: str = port
		self.n0: int = int(n0)
		self.n1: int = int(n1)
		self.exit = self.__exit__
		self.connected: bool = False
		self.f = open(logFile, 'w')
		self.key_left = -2
		self.key_right = -2
	def connect(self):
		"""Connects to the comparator. Waits until connection is established, therefore this method is blocking."""
		s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		print("waiting for connection")
		self.title.configure(text="Waiting for Connection...")
		root.update()
		while True:
			try:
				s.connect((self.ip, int(self.port)))
				break
			except ConnectionRefusedError:
				time.sleep(0.1)
		data: bytes = s.recv(10)
		if data == b"I'm ready!":
			s.send(b'\x02' + self.n0.to_bytes(4, 'little') + self.n1.to_bytes(4, 'little') + b'\x03')
		print("connection established")
		self.title.configure(text="Choose Most Likely to have Signal")
		self.s = s
		self.connected: bool = True
	def __enter__(self):
		"""Sets up the images needed for the study."""
		# get pics, use sorted() to ensure it's the same ones every time
		offset: int = 0
		posNames: list = [self.posDir + img for img in sorted(os.listdir(self.posDir))][offset:offset + self.n1]
		ansNames: list = [self.ansDir + img for img in sorted(os.listdir(self.ansDir))][offset:offset + self.n1]
		negNames: list = [self.negDir + img for img in sorted(os.listdir(self.negDir))][offset:offset + self.n0]
		self.n0: int = len(negNames)
		self.n1: int = len(posNames)
		names: list = negNames + posNames
		with open("names.txt", "w") as f:
			for name in names:
				# write all the files to the file, separated with the file separator
				f.write(name + '\x1C')
		self.images: list = [Image.open(name) for name in names]
		for img in self.images:
			img.thumbnail((IMGWIDTH, IMGHEIGHT))
		self.images: list = [ImageTk.PhotoImage(img) for img in self.images]
		for i, image in enumerate(self.images):
			setattr(image, "filename", names[i])
		ansImages: list = [Image.open(name) for name in ansNames]
		for img in ansImages:
			img.thumbnail((IMGWIDTH, IMGHEIGHT))
		ansImages: list = [ImageTk.PhotoImage(img) for img in ansImages]
		self.ansPairs: list = list(zip(self.images[self.n0:], ansImages))
		self.font: tuple = ('Arial', 72)
		self.correct = ImageTk.PhotoImage(Image.open("correct.png"))
		self.wrong = ImageTk.PhotoImage(Image.open("wrong.png"))
		return self
	def showPics(self, pic1, pic2: int=None):
		"""show the pictures identified by the image number arguments.
		
		If pic2 is not provided, pic1 is a boolean value for displaying either correct or incorrect."""
		if self.mode.get() == 'answers':
			self.img1.configure(image=self.ansPairs[pic1][0])
			self.img2.configure(image=self.ansPairs[pic1][1])
		elif self.mode.get() == 'training' and not isinstance(pic2, int):
			# pic1 is a boolean
				if pic1: #correct
					self.img1.configure(image=self.correct)
					self.img2.configure(image=self.correct)
				else:
					self.img1.configure(image=self.wrong)
					self.img2.configure(image=self.wrong)
		else:
			self.img1.configure(image=self.images[pic1])
			self.img2.configure(image=self.images[pic2])
		root.update()
	def switchModes(self):
		"""Call this when switching modes to make sure the class is ready to start a mode."""
		self.ready: bool = True
		self.decision: int = -1
		self.imgIndex: int = 0
	def run(self):
		"""Finite state machine based on the mode, the connection status, and whether the reader has decided an image or not.
		
		This will call itself on its own, so just make sure to call this method once with <root>.after()"""
		root.update()
		if self.mode.get() == 'training':
			if self.counter <= 0 and self.decision == -1: #put up new pictures
				if self.ready:
					if np.random.random() > 0.5:
						self.imgID1: int = np.random.randint(0, self.n0 - 1)
						self.imgID2: int = np.random.randint(self.n0, self.n1 + self.n0 - 1)
						self.showPics(self.imgID1, self.imgID2)
						self.answer: str = 'right'
					else:
						self.imgID1: int = np.random.randint(self.n0, self.n1 + self.n0 - 1)
						self.imgID2: int =  np.random.randint(0, self.n0 - 1)
						self.showPics(self.imgID1, self.imgID2)
						self.answer: str = 'left'
					self.ready: bool = False
					self.decision: int = -1
				else: # default place, pseudo main loop
					self.showPics(self.imgID1, self.imgID2)
			elif not self.ready and (self.decision == self.key_left or self.decision == self.key_right):
				self.counter: int = 1000
				self.showPics((self.decision == self.key_right and self.answer == 'right') or (self.decision == self.key_left and self.answer == 'left'))
				self.decision: int = -1
				self.ready: bool = True
			elif self.counter > 0:
				self.counter -= 16
				self.ready: bool = self.counter <= 0
		elif self.mode.get() == 'study':
			if not self.connected:
				self.connect()
			if self.ready: #ready for another comparison
				try:
					#print("requesting pics")
					self.s.send(b"send pics!")
					data: bytes = self.s.recv(10)
					#print(data)
					if data:
						if data == b"I'm going!":
							print("Client exited succesfully")
							root.destroy()
						elif data[0] == 2 and data[9] == 3: #valid frame
							self.imgID1: int = int.from_bytes(data[1:5], 'little')
							self.imgID2: int = int.from_bytes(data[5:9], 'little')
							self.data: bytes = data
							#print(img1, img2)
							self.showPics(self.imgID1, self.imgID2)
							self.t1: float = time.time()
							root.update()
							self.ready: bool = False
				except ConnectionResetError:
					print("Client Disconnected")
					self.exit()
				except BrokenPipeError:
					print("Client Disconnected")
					self.exit()
			if self.decision == self.key_left or self.decision == self.key_right:
				if self.decision == self.key_right: #right key
					payload: bytes = (1).to_bytes(4, 'little') + self.data[5:9]
					des: int = self.imgID2
				elif self.decision == self.key_left: #left key
					payload: bytes = (0).to_bytes(4, 'little') + self.data[1:5]
					des: int = self.imgID2
				try:
					self.f.write(f"{self.images[self.imgID1].filename} {self.images[self.imgID2].filename} {self.images[des].filename} {time.time() - self.t1}\n")
					self.s.send(b'\x02' + payload + b'\x03')
				except ConnectionResetError:
					print("Client Disconnected")
					self.exit()
				except BrokenPipeError:
					print("Client Disconnected")
					self.exit()
				self.decision: int = -1
				self.ready: bool = True
		elif self.mode.get() == "answers":
			if self.ready: #put up new pictures
				self.showPics(self.imgIndex)
				self.ready: bool = False
				root.update()
			elif self.decision == self.key_right:
				#print("saw the event")
				self.imgIndex += 1 if self.imgIndex < self.n1 - 1 else 0
				self.decision: int = -1
				self.ready: bool = True
			elif self.decision == self.key_left:
				self.imgIndex -= 1 if self.imgIndex > 0 else 0
				self.decision: int = -1
				self.ready: bool = True
		else:
			if self.counter == 0:
				self.img1.configure(text="Press the key you want to indicate left")
			elif self.counter == 1:
				self.img1.configure(text="Press the key you want to indicate right")
			else:
				self.img1.configure(text="Click the mode you want")
			label.grid(row=0, column=0)

		root.after(16, self.run)
	def pressed(self, event: Event):
		"""Call this when the user has pressed a keyboard button."""
		#print("event!", event.keycode)
		if self.key_left == -2:
			self.key_left = event.keycode
			self.counter += 1
			return
		if self.key_right == -2:
			self.key_right = event.keycode
			self.counter += 1
			return
		self.counter: int = 0
		if event.keycode == self.key_left or event.keycode == self.key_right:
			self.decision: int = event.keycode
	def clicked(self, event: Event):
		"""Call this when the user has clicked the mouse."""
		self.counter: int = 0
		if event.y > 100:
			if event.x > WIDTH - IMGWIDTH:
				self.decision: int = self.key_right
			elif event.x < IMGWIDTH:
				self.decision: int = self.key_left
	def __exit__(self, *args):
		"""Shuts down the connection and closes the window."""
		self.f.close()
		if self.connected:
			self.s.close()
		try:
			root.destroy()
		except Exception: #root already destroyed
			self.counter: int = 0 # to sate the linters
if __name__ == "__main__":
	if len(argv) < 9:	
		print(f"Usage: {__file__} [target present directory] [target absent directory] [answers directory] [merge ip] [merger port] [n0] [n1] [log file]")
	else:
		for i in range(1, 4):
			if not argv[i].endswith('/'):
				argv[i] += '/'
		HEIGHT: int = 700
		WIDTH: int = 1300
		IMGWIDTH: int = 600
		IMGHEIGHT: int = 600
		root = Tk()
		#canvas = Canvas(root, width=WIDTH, height=HEIGHT)
		#canvas.pack()
		#img = Image.open("repository-pic.png")
		#img.thumbnail((IMGWIDTH, IMGHEIGHT))
		#img = ImageTk.PhotoImage(img)
		label: Label = Label(root, text="Choose Most Likely to have Signal", font=('Arial', 56))
		label.grid(row=0, column=0)
		frame: Frame = Frame(root)
		img1: Label = Label(frame)#, image=img)
		img2: Label = Label(frame)#, image=img)
		img1.grid(row=0, column=0, sticky=E)
		img2.grid(row=0, column=2, sticky=W)
		modes: list = [
			("Answers", "answers"),
			("AFC Training", "training"),
			("Study", "study")
		]
		frame.grid_columnconfigure(1, weight=1)
		with AFC(*argv[1:]) as afc:
			afc.title = label
			afc.mode = StringVar()
			afc.mode.set("none")
			afc.img1 = img1
			afc.img2 = img2
			buttons: Frame = Frame(frame)
			for i, (text, mode) in enumerate(modes):
				b: Radiobutton = Radiobutton(buttons, text=text, value=mode, variable=afc.mode, indicatoron=0, command=afc.switchModes)
				b.grid(row=i, column=0)
			buttons.grid(row=0, column=1)
			root.bind("<Key>", afc.pressed)
			root.bind("<Button-1>", afc.clicked)
			root.protocol("WM_DELETE_WINDOW", afc.exit)
			root.after(1, afc.run)
			frame.grid(row=1, column=0)
			mainloop()