#
# Copyright (c) Madrid 2008
# BIT, ETSI Telecomunicacion, UPM
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
#
#
#
import wx
from lxml import etree
from string import atof
GlobalColors = ["white", "red", "blue", "green", "yellow", "cyan" ]
class Section:
def init_from_node(self, node):
if node.tag != "section":
raise ValueError("expected 'section' got '%s'" % (node.tag))
self.points = []
self.color = node.get("color")
for n in node:
if (n.tag == 'point'):
x = atof(n.get("x"))
y = atof(n.get("y"))
self.points.append( (x,y) )
def __init__(self, color=None, points=None, node=None):
self.thickness = 1
if node:
self.init_from_node(node)
else:
self.init_from_color_points(color,points)
def init_from_color_points(self, color, points):
if color:
self.color = color
else:
self.color = GlobalColors[0]
if points:
self.points = points
else:
self.points = []
def Draw(self, dc):
pen = wx.Pen(self.color, self.thickness, wx.SOLID)
brush = wx.Brush(self.color, wx.CROSS_HATCH)
dc.SetPen(pen)
dc.SetBrush(brush)
dc.DrawPolygon(self.points)
def Invert(self):
self.points.reverse()
def Save(self, frame):
section = etree.Element("section", color=self.color)
frame.append(section)
for p in self.points:
cx = "%f" % (p[0])
cy = "%f" % (p[1])
section.append(etree.Element("point", x=cx, y=cy))
def SetPoints(self, points, scale):
self.points = []
for p in points:
self.points.append( ( p[0] * scale[0], p[1] * scale[1]) )
def GetPoints(self, scale):
points = []
for p in self.points:
points.append( ( p[0] * scale[0], p[1] * scale[1]) )
return points
def __eq__(self, other):
if (self.color != other.color):
return False
if len(self.points) != len(other.points):
return False
return self.points == other.points
class ImageSegmentation:
def __init__(self, image=None, sections=None, node=None):
if image:
self.init_from_image_and_sections(image, sections)
elif node:
self.init_from_node(node)
else:
raise ValueError("No image file or node given")
self.image = wx.Image(self.image_file, type=wx.BITMAP_TYPE_ANY, index = -1)
def init_from_image_and_sections(self, image, sections):
self.image_file = image
if sections:
self.Sections = sections
else:
self.Sections = []
def init_from_node(self, node):
if node.tag != "frame":
raise ValueError("Expected 'frame' got '%s' instead" % node.tag)
self.image_file = node.get("image")
self.Sections = []
for n in node:
if n.tag == "section":
self.Sections.append(Section(node=n))
def GetImage(self):
return self.image
def Save(self, root):
frame = etree.Element("frame", image=self.image_file)
if len(self.Sections) > 0:
for sections in self.Sections:
sections.Save( frame )
else:
frame.append(etree.Element("placeholder"))
root.append(frame)
def SetSections(self, sections, out_size):
scale = ( (1.0 * self.image.GetWidth() / out_size[0]),
(1.0 * self.image.GetHeight() / out_size[1]));
self.Sections = []
for s in sections:
self.Sections.append(Section(s.color, s.GetPoints(scale)))
def GetSections(self, out_size):
scale = ( (1.0 * out_size[0]) / self.image.GetWidth(),
(1.0 * out_size[1]) / self.image.GetHeight());
sections = []
for s in self.Sections:
sections.append(Section(s.color, s.GetPoints(scale)))
return sections
def __eq__(self, other):
if len(self.Sections) != len(other.Sections):
return False
for s in self.Sections:
if not (s in other.Sections):
return False
return self.image_file == other.image_file
class WorkSet:
def __init__(self, images=None, node=None):
if node:
self.load(node)
else:
self.init(images)
def load(self, workset):
self.frames = []
if workset.tag != "workset":
raise ValueError("Expected 'workset' got '%s'" % workset.tag)
for f in workset:
self.frames.append( ImageSegmentation( node=f ) )
def init(self, images):
self.frames = []
if not images:
raise ValueError("No images given")
for i in images:
self.frames.append(ImageSegmentation(image=i))
def GetFrameNumber(self):
return len(self.frames)
def GetFrame(self, i):
if i > self.GetFrameNumber():
raise ValueError("Requested frame out of range")
return self.frames[i]
def Save(self):
root = etree.Element("workset")
for f in self.frames:
f.Save(root)
return etree.tostring(root, pretty_print=True)