Added Hough identification

process.py

@@ -7,6 +7,7 @@ from stvid.stars import generate_star_catalog
 from stvid.astrometry import calibrate_from_reference
 from stvid.satellite import generate_satellite_predictions
 from stvid.satellite import find_hough3d_lines
+from stvid.extract import extract_tracks
 import astropy.units as u
 from astropy.utils.exceptions import AstropyWarning
 from astropy.coordinates import EarthLocation
@@ -14,9 +15,9 @@ import warnings
 import configparser
 import argparse
 import os
+from termcolor import colored
 
 if __name__ == "__main__":
-
     # Read commandline options
     conf_parser = argparse.ArgumentParser(description='Process captured' +
                                                       ' video frames.')
@@ -51,14 +52,36 @@ if __name__ == "__main__":
     os.chdir(args.file_dir)
 
     # Get files
-    files = sorted(glob.glob("2*.fits"))
+    fnames = sorted(glob.glob("2*.fits"))
 
     # Statistics file
     fstat = open("imgstat.csv", "w")
     fstat.write("fname,mjd,ra,de,rmsx,rmsy,mean,std,nstars,nused\n")
 
+    # Extract settings
+    # Minimum predicted velocity (pixels/s)
+    drdtmin = 10.0
+
+    # Track selection region around prediction (pixels)
+    trkrmin = 10.0
+
+    # Track selection sigma
+    trksig = 5.0
+
+    # Minimum track points
+    ntrkmin = 10
+
+    # Create output dirs
+    path = args.file_dir
+    if not os.path.exists(os.path.join(path, "classfd")):
+        os.makedirs(os.path.join(path, "classfd"))
+    if not os.path.exists(os.path.join(path, "catalog")):
+        os.makedirs(os.path.join(path, "catalog"))
+    if not os.path.exists(os.path.join(path, "unid")):
+        os.makedirs(os.path.join(path, "unid"))
+
     # Loop over files
-    for fname in files:
+    for fname in fnames:
         # Generate star catalog
         pix_catalog = generate_star_catalog(fname)
 
@@ -69,25 +92,26 @@ if __name__ == "__main__":
         # Generate satellite predictions
         generate_satellite_predictions(fname)
 
-        # Extract lines with 3D Hough transform
+        # Detect lines with 3D Hough transform
         ids = find_hough3d_lines(fname)
 
+        # Extract tracks
+        extract_tracks(fname, trkrmin, drdtmin, trksig, ntrkmin, path)
+        
         # Stars available and used
         nused = np.sum(pix_catalog.flag == 1)
         nstars = pix_catalog.nstars
-
+        
         # Get properties
         ff = fourframe(fname)
 
-        print(("%s,%.8lf,%.6f,%.6f,%.3f,%.3f," +
-              "%.3f,%.3f,%d,%d") % (ff.fname, ff.mjd, ff.crval[0],
-                                    ff.crval[1], 3600*ff.crres[0],
-                                    3600*ff.crres[1], np.mean(ff.zavg),
-                                    np.std(ff.zavg), nstars, nused))
+        # Write output
+        output = "%s %10.6f %10.6f %4d/%4d %5.1f %5.1f %6.2f +- %6.2f"%(ff.fname, ff.crval[0], ff.crval[1], nused, nstars, 3600.0*ff.crres[0], 3600.0*ff.crres[1], np.mean(ff.zavg), np.std(ff.zavg))
+        print(colored(output, "yellow"))
         fstat.write(("%s,%.8lf,%.6f,%.6f,%.3f,%.3f,%.3f," +
                      "%.3f,%d,%d\n") % (ff.fname, ff.mjd, ff.crval[0],
                                         ff.crval[1], 3600*ff.crres[0],
                                         3600*ff.crres[1], np.mean(ff.zavg),
                                         np.std(ff.zavg), nstars, nused))
-
     fstat.close()
+