1#!/usr/local/bin/python3.8
2# This file is part of the Astrometry.net suite.
3# Licensed under a 3-clause BSD style license - see LICENSE
4from __future__ import print_function
5
6import sys
7try:
8    import pyfits
9except ImportError:
10    try:
11        from astropy.io import fits as pyfits
12    except ImportError:
13        raise ImportError("Cannot import either pyfits or astropy.io.fits")
14
15from math import *
16from numpy import *
17from pylab import *
18from scipy.ndimage.filters import *
19from astrometry.util.fits import pyfits_writeto
20
21def normalized_hough(x, y, imgw, imgh, rlo, rhi, tlo, thi, nr, nt):
22    houghimg = zeros((nr, nt)).astype(int)
23    tstep = (thi - tlo) / float(nt)
24    rstep = (rhi - rlo) / float(nr)
25    # For each point, accumulate into the Hough transform image...
26    tt = tlo + (arange(nt) + 0.5) * tstep
27    cost = cos(tt)
28    sint = sin(tt)
29    for (xi,yi) in zip(x, y):
30        rr = xi * cost + yi * sint
31        ri = floor((rr - rlo) / rstep).astype(int)
32        I = (ri >= 0) * (ri < nr)
33        houghimg[ri[I], I] += 1
34
35    # Compute the approximate Hough normalization image
36    houghnorm = zeros((nr, nt)).astype(float)
37    rr = rlo + (arange(nr) + 0.5) * rstep
38    for ti,t in enumerate(tt):
39        (x0,x1,y0,y1) = clip_to_image(rr, t, imgw, imgh)
40        dist = sqrt((x0 - x1)**2 + (y0 - y1)**2)
41        houghnorm[:, ti] = dist
42    # expected number of points: dist is the length of the slice,
43    #   rstep is the width of the slice; len(x)/A is the source density.
44    houghnorm *= (rstep * len(x) / (imgw*imgh))
45    return (houghimg, houghnorm, rr, tt, rstep, tstep)
46
47def clip_to_image(r, t, imgw, imgh):
48    eps = 1e-9
49    if abs(t) < eps or abs(t-pi) < eps:
50        # near-vertical.
51        s = (abs(t) < eps) and 1 or -1
52        y0 = 0
53        y1 = ((r*s >= 0) * (r*s < imgw)) * imgh
54        x0 = x1 = clip(r, 0, imgw)
55        return (x0, x1, y0, y1)
56    m = -cos(t)/sin(t)
57    b = r/sin(t)
58    x0 = 0
59    x1 = imgw
60    y0 = clip(b + m*x0, 0, imgh)
61    y1 = clip(b + m*x1, 0, imgh)
62    x0 = clip((y0 - b) / m, 0, imgw)
63    x1 = clip((y1 - b) / m, 0, imgw)
64    y0 = clip(b + m*x0, 0, imgh)
65    y1 = clip(b + m*x1, 0, imgh)
66    return (x0, x1, y0, y1)
67
68def removelines_general(infile, outfile, nt=180, nr=180, thresh1=2.,
69                        thresh2=5., plots=False):
70    p = pyfits.open(infile)
71    xy = p[1].data
72    hdr = p[1].header
73    x = xy.field('X').copy()
74    y = xy.field('Y').copy()
75
76    imshowargs = { 'interpolation':'nearest', 'origin':'lower' }
77
78    imgw = int(ceil(max(x) - min(x)))
79    imgh = int(ceil(max(y) - min(y)))
80
81    x -= min(x)
82    y -= min(y)
83
84    Rmax = sqrt(imgw**2 + imgh**2)
85    Rmin = -Rmax
86
87    (houghimg, houghnorm, rr, tt, rstep, tstep
88     ) = normalized_hough(x, y, imgw, imgh, Rmin, Rmax, 0, pi, nr, nt)
89
90    hnorm = houghimg / maximum(houghnorm, 1)
91
92    if plots:
93        clf()
94        plot(x,y,'r.')
95        savefig('xy.png')
96
97        clf()
98        imshow(houghimg, **imshowargs)
99        xlabel('Theta')
100        ylabel('Radius')
101        colorbar()
102        savefig('hough.png')
103
104        clf()
105        imshow(houghnorm, **imshowargs)
106        xlabel('Theta')
107        ylabel('Radius')
108        colorbar()
109        savefig('norm.png')
110
111        clf()
112        imshow(hnorm, **imshowargs)
113        xlabel('Theta')
114        ylabel('Radius')
115        colorbar()
116        savefig('hnorm.png')
117
118
119    I = find(hnorm.ravel() >= thresh1)
120    print('%i peaks are above the coarse threshold' % len(I))
121    bestri = I / nt
122    bestti = I % nt
123
124    if plots:
125        a=axis()
126        for (ri,ti) in zip(bestri,bestti):
127            plot([ti-2, ti-2, ti+2, ti+2, ti-2], [ri-2, ri+2, ri+2, ri-2, ri-2], 'r-')
128            axis(a)
129            savefig('zooms.png')
130
131        clf()
132        plot(x,y,'r.')
133        for (ri,ti) in zip(bestri,bestti):
134            (x0,x1,y0,y1) = clip_to_image(rr[ri], tt[ti], imgw, imgh)
135            plot([x0,x1],[y0,y1], 'b-')
136        savefig('xy2.png')
137
138    # how big a search area around each peak?
139    boxsize = 1
140    # how much more finely to grid.
141    finer = 3
142    nr2 = (boxsize * 2)*finer + 1
143    nt2 = nr2
144
145    bestrt = []
146    keep = array([True] * len(x))
147    for (ri,ti) in zip(bestri,bestti):
148        (subh, subnorm, subrr, subtt, subrstep, subtstep
149         ) = normalized_hough(x, y, imgw, imgh,
150                              rr[max(ri-boxsize, 0)], rr[min(ri+boxsize, nr-1)],
151                              tt[max(ti-boxsize, 0)], tt[min(ti+boxsize, nt-1)],
152                              nr2, nt2)
153        #print '  median normalization:', median(subnorm)
154        subhnorm = subh / maximum(subnorm,1)
155        I = find((subhnorm).ravel() >= thresh2)
156        for i in I:
157            bestsubri = i / nt2
158            bestsubti = i % nt2
159            r = subrr[bestsubri]
160            t = subtt[bestsubti]
161            bestrt.append((r,t))
162            #print '  (r=%.1f, t=%.1f): factor %.1f above expected' % (r, t*180/pi, subhnorm.ravel()[i])
163            thisr = x * cos(t) + y * sin(t)
164            keep *= (abs(thisr - r) > subrstep/2.)
165
166    print('In finer grid: found %i peaks' % len(bestrt))
167
168    if plots:
169        clf()
170        subplot(1,1,1)
171        plot(x,y,'r.')
172        for (r,t) in bestrt:
173            (x0,x1,y0,y1) =  clip_to_image(r, t, imgw, imgh)
174            plot([x0,x1],[y0,y1],'b-')
175        savefig('xy3.png')
176
177        clf()
178        plot(x,y,'r.')
179        plot(x[keep == False], y[keep == False], 'b.')
180        savefig('xy4.png')
181
182    p[1].data = p[1].data[keep]
183    pyfits_writeto(p, outfile)
184    return 0
185
186def exact_hough_normalization():
187    houghnorm = zeros((nr, nt)).astype(float)
188    [xx,yy] = meshgrid(arange(imgw), arange(imgh))
189    yyflat = yy.ravel()
190    xxflat = xx.ravel()
191    for ti in range(nt):
192        print(ti)
193        t = (ti+0.5) * tstep
194        rr = xxflat * cos(t) + yyflat * sin(t)
195        ri = floor((rr - Rmin) / rstep).astype(int)
196        (counts, nil) = histogram(ri, range(0, nr+1))
197        houghnorm[:, ti] += counts
198    clf()
199    imshow(houghnorm, **imshowargs)
200    colorbar()
201    savefig('houghnorm.png')
202
203
204
205if __name__ == '__main__':
206    args = sys.argv[1:]
207    plots = False
208    if '-p' in args:
209        plots = True
210        args.remove('-p')
211
212    if len(args) != 2:
213        print('Usage: %s [options] <input-file> <output-file>' % sys.argv[0])
214        print('   [-p]: create plots')
215        exit(-1)
216
217    infile = args[0]
218    outfile = args[1]
219    rtncode = removelines_general(infile, outfile, plots=plots)
220    sys.exit(rtncode)
221
222