freedraw.py

import sys
import numpy as np
from PySide import QtCore, QtGui

class DrawWindow(QtGui.QWidget):
    def __init__(self, parent):
        QtGui.QWidget.__init__(self, parent)
        #constants
        self.PEN_WEIGHT = 8
        self.pen_color = QtCore.Qt.cyan
        self.pen_color_long = 4278255615


        self.first_point = QtCore.QPoint()
        self.intermediate_point = QtCore.QPoint()

        self.left_point = None
        self.right_point = None
        self.top_point = None
        self.bottom_point = None

        self.drawing = False

        self.image = QtGui.QImage(QtCore.QSize(500, 500), QtGui.QImage.Format_RGB32)
        self.image.fill(QtGui.qRgb(255, 255, 255))

        self.fill_matrix = np.zeros((500, 500), dtype=int)
    
    def isValidPoint(self, point):
        return(point.x() <= 500 and point.x() >= 0 and point.y() <= 500 and point.y() >= 0)

    def findCenter(self):
        # Check for 2 intersections to left of x or right of x.  Do the same for y.
        # If there are 2 intersections for any of the 4 directions, this means the point is not enclosed in the curve.

        starting_point = ((self.left_point.x() + self.right_point.x())/2, (self.top_point.y() + self.bottom_point.y())/2)

        painter = QtGui.QPainter(self.image)
        painter.setPen(QtGui.QPen(QtCore.Qt.red, self.PEN_WEIGHT, QtCore.Qt.SolidLine, QtCore.Qt.RoundCap, QtCore.Qt.RoundJoin))


        starting_x, starting_y = starting_point[0], starting_point[1]

        painter.drawRect(starting_x, starting_y, 1, 1)

        # increment by pen weight in order to avoid directly adjacent points with same color
        intersections = []
        x = 0
        while(x < starting_x):
            if(self.image.pixel(x, starting_y) == self.pen_color_long):
                intersections.append(x)
                x += self.PEN_WEIGHT*2
            else:
                x+=1
        num_intersections = len(intersections)
        if(num_intersections > 0 and num_intersections % 2 == 0):
            # return average of 2 intersection points as new center
            return((intersections[0] + intersections[1])/2, starting_y)
        
        intersections = []
        x = starting_x
        while(x < 500):
            if(self.image.pixel(x, starting_y) == self.pen_color_long):
                intersections.append(x)
                x += self.PEN_WEIGHT*2
            else:
                x += 1
        num_intersections = len(intersections)
        if(num_intersections > 0 and num_intersections % 2 == 0):
            # return average of 2 intersection points as new center
            return((intersections[0] + intersections[1])/2, starting_y)

        return(starting_point)

    
    def savePixels(self):
        for x in range(500):
            for y in range(500):
                if(self.image.pixel(x,y) == self.pen_color_long):
                    self.fill_matrix[x][y] = 1
        
        np.savetxt("data.csv", self.fill_matrix, delimiter=",", fmt="%u")


    def mousePressEvent(self, event):
        current_point = event.pos()
        if(self.isValidPoint(current_point)):
            self.first_point = self.intermediate_point = current_point
            self.drawing = True
    
    def mouseMoveEvent(self, event):
        current_point = event.pos()

        if(self.drawing and event.buttons() and self.isValidPoint(current_point)):
            self.draw_line_to(self.intermediate_point, current_point)        
            self.intermediate_point = current_point
        
        if(not self.left_point or current_point.x() < self.left_point.x()):
            self.left_point = current_point
        if(not self.right_point or current_point.x() > self.right_point.x()):
            self.right_point = current_point
        if(not self.top_point or current_point.y() < self.top_point.y()):
            self.top_point = current_point
        if(not self.bottom_point or current_point.y() > self.bottom_point.y()):
            self.bottom_point = current_point

    def mouseReleaseEvent(self, event):
        if(self.drawing and event.buttons()):
            self.drawing = False
        
        self.draw_line_to(self.first_point, event.pos())
        self.fill_curve()
    
    def paintEvent(self, event):
        painter = QtGui.QPainter(self)
        painter.drawImage(event.rect(), self.image)

    def draw_line_to(self, startpoint, endpoint):
        painter = QtGui.QPainter(self.image)
        painter.setRenderHint(QtGui.QPainter.Antialiasing)
        painter.setPen(QtGui.QPen(self.pen_color, self.PEN_WEIGHT, QtCore.Qt.SolidLine, QtCore.Qt.RoundCap, QtCore.Qt.RoundJoin))
        painter.drawLine(startpoint, endpoint)

        self.update()

    def fill_curve(self):
        to_visit = [self.findCenter()]
        self.fill_matrix[to_visit[0][0]][to_visit[0][1]] = 1

        painter = QtGui.QPainter(self.image)
        painter.setPen(QtGui.QPen(self.pen_color, 1, QtCore.Qt.SolidLine, QtCore.Qt.RoundCap, QtCore.Qt.RoundJoin))

        while(len(to_visit) > 0):
            x,y = to_visit.pop()
            if(x < 0 or y < 0):
                break

            # For x+2 and y+2, +2 is used because painter.drawRect() covers up the +1 coordinate (I think)
            if(self.image.pixel(x,y+2) != self.pen_color_long):
                to_visit.append((x,y+2))
            if(self.image.pixel(x+2,y) != self.pen_color_long):
                to_visit.append((x+2,y))
            if(self.image.pixel(x,y-1) != self.pen_color_long):
                to_visit.append((x,y-1))
            if(self.image.pixel(x-1,y) != self.pen_color_long):
                to_visit.append((x-1,y))
            painter.drawRect(x,y,1,1)

        self.update()
  
        self.left_point = self.right_point = self.top_point = self.bottom_point = None

        self.savePixels()

#Note: not working yet
class ButtonBar(QtGui.QWidget):
    def __init__(self, parent):
        QtGui.QWidget.__init__(self, parent)

        self.layout = QtGui.QHBoxLayout()
        self.save_button = QtGui.QPushButton("Save")
        self.reset_button = QtGui.QPushButton("Reset")
        self.layout.addWidget(self.save_button, 0, QtCore.Qt.Alignment())
        self.layout.addWidget(self.reset_button, 0, QtCore.Qt.Alignment())


class MainWindow(QtGui.QWidget):
    def __init__(self):
        app = QtGui.QApplication(sys.argv)
        QtGui.QWidget.__init__(self)
        
        self.resize(500,500)
        self.layout = QtGui.QVBoxLayout()

        self.draw_window = DrawWindow(self)
        self.button_bar = ButtonBar(self)

        self.layout.addWidget(self.draw_window, 0, QtCore.Qt.Alignment())
        self.layout.addWidget(self.button_bar, 0, QtCore.Qt.Alignment())
        self.setLayout(self.layout)

        self.show()

        app.exec_()

window = MainWindow()