Vector.py 2.1 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  1. try:
  2. import cPickle as pickle
  3. except ImportError:
  4. import pickle
  5. import sample
  6. class FeatureVector:
  7. def __init__(self):
  8. self.features = []
  9. self.activefeatures = []
  10. # list of key, value tuples represnting values for features.
  11. self.sampleInfo = {}
  12. self.classification = None
  13. # set which features are active using a binary list
  14. # ex: For 2nd, 4th and 5th features to be active, pass [0,1,0,1,1]
  15. def setFeatures(self, binlist):
  16. activefeatures = []
  17. if len(binlist) != len(self.features):
  18. print("Feature choice list must equal length of feature list")
  19. for b in range(len(binlist)):
  20. if binlist[b] == 1:
  21. activefeatures.append(self.features[b])
  22. self.activefeatures = activefeatures
  23. def __repr__(self):
  24. return str(self.features)
  25. # use to make a sample into a vector
  26. def SampleToFeatureVector(s: sample):
  27. print(s.__activities)
  28. v = FeatureVector()
  29. for a in s.__activities.keys():
  30. for d in s.__activities[a].keys():
  31. # a + d looks like: high.total_packets
  32. v.sampleInfo.update({a + "." + d: s.__activities[a][d]})
  33. v.features.extend(a["high"].values())
  34. v.features.extend(a["mid"].values())
  35. v.features.extend(a["low"].values())
  36. v.sampleInfo.update(("total_time", s.total_time))
  37. v.sampleInfo.update(("average_iat", s.average_iat))
  38. v.sampleInfo.update(("dead_time", s.dead_time))
  39. v.features.append(s.total_time)
  40. v.features.append(s.average_iat)
  41. v.features.append(s.dead_time)
  42. return v
  43. def writePickledData(filename):
  44. v = FeatureVector()
  45. v.features = [1, 2, 3, 4]
  46. v.activefeatures = [1, 2, 3, 4]
  47. vs = [v,v]
  48. with open(filename, 'wb') as file:
  49. pickle.dump(vs, file)
  50. def readPickledData(filename):
  51. with open(filename, 'rb') as file:
  52. x = pickle.load(file)
  53. # print(x)
  54. return x
  55. def main():
  56. fv = FeatureVector()
  57. writePickledData("test.bin")
  58. readPickledData("test.txt")
  59. s = sample.Sample(3)
  60. fv = SampleToFeatureVector()
  61. if __name__ == '__main__':
  62. main()