research/TCS/pressure/get_data.py~

   1 #! /usr/bin/python
   2
   3 from PIL import Image
   4 import Image, ImageDraw
   5 import sys
   6 import numpy
   7
   8
   9 def PixelGood(pixel, threshold):
  10         #print str(pixel)
  11         return GreyScale(pixel) < threshold;
  12         #return (pixel[1] > (pixel[0] + pixel[2]) and pixel[1] > threshold)
  13
  14 def rect_size(e):
  15         return (e[3] - e[1]) * (e[2] - e[0])
  16
  17 def GreyScale(pixel):
  18         greyscale = 0
  19         for i in range(0, len(pixel)):
  20                 greyscale += pixel[i]
  21         greyscale = greyscale / len(pixel)
  22         return greyscale
  23
  24 def SubDivide(pix, bounds, fill_fraction=0.1, threshold=200, seperation=5, min_area=1000):
  25         rect = []
  26
  27         xStart = bounds[2]
  28         xEnd = bounds[0]
  29         yStart = bounds[3]
  30         yEnd = bounds[1]
  31         xLast = bounds[0]
  32         good = 0
  33         for x in range(bounds[0], bounds[2]):
  34                 for y in range(bounds[1], bounds[3]):
  35                         if PixelGood(pix[x, y], threshold):
  36                                 if (x - xLast) > seperation:
  37                                         if ((xEnd - xStart) * (yEnd - yStart) > min_area):
  38                                                 rect.append([xStart, yStart, xEnd, yEnd])
  39
  40                                         xStart = bounds[2]
  41                                         xEnd = bounds[0]
  42                                         yStart = bounds[3]
  43                                         yEnd = bounds[1]
  44                                         good = 0
  45
  46                                 if (x < xStart):
  47                                         xStart = x
  48                                 if (y < yStart):
  49                                         yStart = y
  50                                 if (x > xEnd):
  51                                         xEnd = x
  52                                 if (y > yEnd):
  53                                         yEnd = y
  54                                 xLast = x
  55                                 good += 1
  56
  57         return rect
  58
  59 def Process(fileName, digits):
  60
  61         results = []
  62         image = Image.open(fileName)
  63         pix = image.load()
  64         rect = SubDivide(pix, (0, 0, image.size[0], image.size[1]))
  65         #print "Found " + str(len(rect)) + " rectangles"
  66         test = Image.new("RGB", image.size, "white")
  67         test.paste(image)
  68         draw1 = ImageDraw.Draw(test)
  69         for i in range(0, len(rect)):
  70                 #draw1.rectangle(rect[i], fill=None, outline="red")
  71
  72
  73                 #img.show()
  74                 closest = PickClosest(pix, rect[i], digits)
  75                 print("Best result is " + str(closest))
  76                 if (closest == None or closest[1] < 0.0):
  77
  78                         img = Image.new("RGB", (rect[i][2]-rect[i][0], rect[i][3]-rect[i][1]), "white")
  79                         draw2 = ImageDraw.Draw(img)
  80                         for x in range(rect[i][0], rect[i][2]):
  81                                 for y in range(rect[i][1], rect[i][3]):
  82                                         draw2.rectangle((x-rect[i][0], y-rect[i][1], x-rect[i][0], y-rect[i][1]), fill=pix[x, y], outline=None)
  83
  84                         img.show()
  85                         sys.stdout.write("Enter digit: ")
  86                         closest = (int(sys.stdin.readline().strip("\r\n ")), 1)
  87                         img.save(str(closest[0])+".png")
  88                 results.append(closest[0])
  89         return results
  90
  91 def PixToList(pix, bounds):
  92         a = []
  93         for x in range(bounds[0], bounds[2]):
  94                 a.append([])
  95
  96         for x in range(bounds[0], bounds[2]):
  97                 for y in range(bounds[1], bounds[3]):
  98
  99                         a[x - bounds[0]].append(GreyScale(pix[x, y]))
 100                         #print(str(x) + "," + str(y), len(fft[x - bounds[0]]))
 101                         #print(str(fft[x - bounds[0]][y - bounds[1]]))
 102         return a
 103
 104 def PickClosest(pix, bounds, digits, boxes):
 105         scores = []
 106
 107
 108
 109
 110
 111         a = numpy.array(PixToList(pix, bounds))
 112         b = abs(numpy.fft.rfft2(a))
 113         #Image.fromarray(b).show()
 114
 115         for d in digits:
 116                 score = 0
 117                 for x in range(0, len(b)):
 118                         for y in range(0, len(b[x])):
 119                                 if (x >= len(d[1]) or y >= len(d[1][x])):
 120                                         score -= 1
 121                                 else:
 122                                         score -= abs(b[x][y] - d[1][x][y])
 123                 scores.append([d[0], float(score)])
 124         return sorted(scores, key = lambda e : e[1], reverse = False)[0]
 125
 126 if __name__ == "__main__":
 127
 128         if (len(sys.argv) != 2):
 129                 print("Usage " + str(sys.argv[0]) + " input_image")
 130
 131         digits = []
 132         for i in range(0, 10):
 133                 try:
 134                         digit = Image.open(str(i)+".png")
 135                 except:
 136                         continue
 137
 138
 139                 digits.append((i, digit.load(), digit.size))
 140
 141
 142         #sys.exit(0)
 143         r = Process(sys.argv[1], digits)
 144         print(str(r[0]) + "." + str(r[1]) + str(r[2]) + "e-"+str(r[3]))
 145
 146
 147         sys.exit(0)
 148
 149
 150