tracking.py 1.6 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. import re
  2. from pathlib import Path
  3. import pandas as pd
  4. import videobatch
  5. # for extraction of ROI names from files
  6. ROI_PATTERN = re.compile(r"_([a-zA-Z]+)\.png")
  7. # base spec, used in videobatch
  8. base_spec = {
  9. "command": "pixylation",
  10. "mode": "CM",
  11. "colors": dict(Magenta=[260,360], Green=[90,225]),
  12. "sources": ["*.mp4"]
  13. }
  14. # the root repository dir
  15. rootdir = Path().resolve().parent
  16. # dataset dirs
  17. rawdir = rootdir / "datasets" / "raw"
  18. trkdir = rootdir / "datasets" / "tracking"
  19. # read info from SESSIONS
  20. sessions = pd.read_csv(str(rootdir / "SESSIONS.csv"))
  21. # iterate over sessions
  22. for rowidx in range(sessions.shape[0]):
  23. session = sessions.iloc[rowidx]
  24. sname = f"session{session.date}-{session.index:03d}"
  25. # source video directory
  26. srcdir = rawdir / session.subject / sname / "video"
  27. # read ROI info
  28. roidir = trkdir / session.subject / sname / "ROI" # may not be in the repository
  29. ROIs = {}
  30. for child in roidir.iterdir():
  31. is_ROIfile = ROI_PATTERN.search(child.name)
  32. if not is_ROIfile:
  33. continue
  34. else:
  35. ROIname = is_ROIfile.group(1)
  36. ROIs[ROIname] = str(child)
  37. # output directories
  38. resultdir = trkdir / session.subject / sname / "tracked"
  39. maskdir = trkdir / session.subject / sname / "mask" # not in the repository
  40. # configure specs
  41. spec = base_spec.copy()
  42. spec['sourcedir'] = str(srcdir)
  43. spec['Rois'] = ROIs
  44. spec['resultdir'] = str(resultdir)
  45. spec['maskdir'] = str(maskdir)
  46. # run videobatch for this session
  47. videobatch.run(spec)