week6/week6.py

   1 #!/usr/bin/env python
   2 '''
   3 AY250 HW6
   4 Image search and manipulation GUI using Traits
   5 Author: Bryan Steinbach
   6 '''
   7
   8
   9 import urllib
  10
  11 import wx
  12 import numpy as np
  13 from enthought.traits.api import HasTraits,Str,Float,Button,Instance
  14 from enthought.traits.ui.api import View,Item,ButtonEditor,VSplit,HSplit
  15 from yahoo.search.image import ImageSearch
  16 from PIL import Image
  17 from mpl_figure_editor import MPLFigureEditor
  18 from matplotlib.figure import Figure
  19 import pylab
  20
  21 def yahoosearch(query):
  22         ''' Yahoo searches for an image and returns a URL, or an empty string if something goes wrong '''
  23         try:
  24                 search = ImageSearch(app_id="junk",query=query)
  25                 search.results = 1
  26                 for res in search.parse_results():
  27                         return res.Url
  28         except:
  29                 print "Image searched failed!"
  30                 pass
  31
  32         return ""
  33
  34 def fetchlink(url):
  35         ''' Fetch file from link and return path to local file '''
  36         return urllib.urlretrieve(url)[0]
  37
  38 yuvmat = np.array([[0.299,0.587,0.114],[-0.14713,-0.28886,0.436],[0.615,-0.51499,-0.10001]])            # source: wikipedia YUV article
  39 yuvmatinv = np.linalg.inv(yuvmat)
  40 print yuvmatinv
  41
  42 class ImgSearch(HasTraits):
  43         # Model
  44         query = Str("Enter search")
  45         run_query = Button()
  46         url = Str
  47
  48         refresh = Button()
  49         rotate90 = Button()
  50         rotate_color = Button()
  51         rgb_to_yuv = Button()
  52         yuv_to_rgb = Button()
  53
  54         query = "red star winter orbit"
  55
  56         figure = Instance(Figure,())
  57
  58         # Controller - button callbacks
  59         def _refresh_fired(self):
  60                 ''' Reload raw image from disk '''
  61                 self.load_image()
  62                 self.draw_image()
  63
  64         def _rotate90_fired(self):
  65                 ''' Rotate image 90 degrees clockwise '''
  66                 self.ima = pylab.rot90(self.ima)
  67                 self.draw_image()
  68
  69         def _rgb_to_yuv_fired(self):
  70                 ''' Show YUV as RGB '''
  71                 self.color_transform(yuvmat)
  72                 self.draw_image()
  73
  74         def _yuv_to_rgb_fired(self):
  75                 ''' Show RGB as YUV '''
  76                 self.color_transform(yuvmatinv)
  77                 self.draw_image()
  78
  79         def _rotate_color_fired(self):
  80                 ''' Advance color space RGB -> BRG '''
  81                 x = np.zeros_like(self.ima)
  82                 r,g,b = self.ima[:,:,0],self.ima[:,:,1],self.ima[:,:,2]
  83                 x[:,:,0] = b
  84                 x[:,:,1] = r
  85                 x[:,:,2] = g
  86                 self.ima = x
  87                 self.draw_image()
  88
  89
  90         def _run_query_fired(self):
  91                 ''' Search for, download and display image from query text'''
  92                 self.fetch_image()
  93                 self.load_image()
  94                 self.draw_image()
  95
  96         # Controller helper methods
  97         def draw_image(self):
  98                 ''' Redraw image on canvas '''
  99                 self.figure.axes[0].images=[]
 100                 self.figure.axes[0].imshow(self.ima,aspect='auto',interpolation='nearest')
 101                 wx.CallAfter(self.figure.canvas.draw)
 102
 103         def load_image(self):
 104                 ''' Load image from disk '''
 105                 im = Image.open(self.localfile)
 106                 self.ima = np.asarray(im)/256.0
 107                 print "ima.shape: ",self.ima.shape
 108
 109
 110         def fetch_image(self):
 111                 ''' Search for and download image '''
 112                 print "query: ",self.query
 113                 self.url = yahoosearch(self.query)
 114                 print "url: ",self.url
 115                 self.localfile = fetchlink(self.url)
 116                 print "localfile: ",self.localfile
 117
 118         def color_transform(self,mat):
 119                 ''' Linear colorspace transform '''
 120                 x = np.zeros_like(self.ima)
 121                 for i in range(self.ima.shape[0]):
 122                         for j in range(self.ima.shape[1]):
 123                                 x[i,j,:] = np.dot(mat,self.ima[i,j,:])
 124                 self.ima = x
 125
 126
 127         # View
 128         qi = Item('query',label='Input',width=250)
 129         rqi = Item('run_query',show_label=False)
 130         figi = Item('figure',editor=MPLFigureEditor(),show_label=False)
 131         urli = Item('url',label='Image URL')
 132
 133         refi = Item('refresh',show_label=False)
 134         rot90i = Item('rotate90',show_label=False)
 135         rgbyuvi = Item('rgb_to_yuv',show_label=False)
 136         yuvrgbi = Item('yuv_to_rgb',show_label=False)
 137         rci = Item('rotate_color',show_label=False)
 138
 139         view = View(HSplit(qi,rqi),urli,figi,HSplit(refi,rot90i,rci,rgbyuvi,yuvrgbi),
 140                         width=400,height=300,resizable=True)
 141
 142         def __init__(self):
 143                 super(ImgSearch,self).__init__()
 144                 self.figure.add_subplot(111)
 145
 146
 147 c = ImgSearch()
 148 c.configure_traits()
 149