Commit bdb92e42ff0bd165dfeb6ca4ef40b9e51c64044b

Authored by tfmoraes
1 parent f635c217

ENH: Using threads to open dicoms

Showing 1 changed file with 49 additions and 13 deletions   Show diff stats
invesalius/reader/dicom_reader.py
@@ -16,18 +16,23 @@ @@ -16,18 +16,23 @@
16 # PARTICULAR. Consulte a Licenca Publica Geral GNU para obter mais 16 # PARTICULAR. Consulte a Licenca Publica Geral GNU para obter mais
17 # detalhes. 17 # detalhes.
18 #-------------------------------------------------------------------------- 18 #--------------------------------------------------------------------------
19 -from vtk.util.colors import yellow  
20 import glob 19 import glob
21 -import os  
22 import math 20 import math
  21 +import os
  22 +import Queue
  23 +import tempfile
  24 +import threading
  25 +
  26 +from multiprocessing import cpu_count
23 27
24 import vtk 28 import vtk
25 import vtkgdcm 29 import vtkgdcm
26 import gdcm 30 import gdcm
27 -import thread  
28 import wx 31 import wx
29 import wx.lib.pubsub as ps 32 import wx.lib.pubsub as ps
30 33
  34 +from vtk.util.colors import yellow
  35 +
31 import constants as const 36 import constants as const
32 import dicom 37 import dicom
33 import dicom_grouper 38 import dicom_grouper
@@ -73,12 +78,38 @@ def SortFiles(filelist, dicom): @@ -73,12 +78,38 @@ def SortFiles(filelist, dicom):
73 78
74 return filelist 79 return filelist
75 80
  81 +class LoadDicom(threading.Thread):
  82 + def __init__(self, grouper, tempdir, q, l):
  83 + threading.Thread.__init__(self)
  84 + self.grouper = grouper
  85 + self.tempdir = tempdir
  86 + self.q = q
  87 + self.l = l
  88 + def run(self):
  89 + grouper = self.grouper
  90 + q = self.q
  91 + tempdir = self.tempdir
  92 + while 1:
  93 + filepath = q.get()
  94 + print "thread %s consumed %s" % (self.getName(), filepath)
  95 + if not filepath:
  96 + break
  97 + parser = dicom.Parser()
  98 + if parser.SetFileName(filepath):
  99 + dcm = dicom.Dicom()
  100 + self.l.acquire()
  101 + dcm.SetParser(parser)
  102 + grouper.AddFile(dcm)
  103 + self.l.release()
  104 + dcm.image.SetTempDir(tempdir)
  105 +
76 106
77 def yGetDicomGroups(directory, recursive=True, gui=True): 107 def yGetDicomGroups(directory, recursive=True, gui=True):
78 """ 108 """
79 Return all full paths to DICOM files inside given directory. 109 Return all full paths to DICOM files inside given directory.
80 """ 110 """
81 nfiles = 0 111 nfiles = 0
  112 + tempdir = tempfile.mkdtemp()
82 # Find total number of files 113 # Find total number of files
83 if recursive: 114 if recursive:
84 for dirpath, dirnames, filenames in os.walk(directory): 115 for dirpath, dirnames, filenames in os.walk(directory):
@@ -89,31 +120,36 @@ def yGetDicomGroups(directory, recursive=True, gui=True): @@ -89,31 +120,36 @@ def yGetDicomGroups(directory, recursive=True, gui=True):
89 120
90 counter = 0 121 counter = 0
91 grouper = dicom_grouper.DicomPatientGrouper() 122 grouper = dicom_grouper.DicomPatientGrouper()
  123 + q = Queue.Queue()
  124 + l = threading.Lock()
  125 + threads = []
  126 + for i in xrange(cpu_count()):
  127 + t = LoadDicom(grouper, tempdir, q, l)
  128 + t.start()
  129 + threads.append(t)
92 # Retrieve only DICOM files, splited into groups 130 # Retrieve only DICOM files, splited into groups
93 if recursive: 131 if recursive:
94 for dirpath, dirnames, filenames in os.walk(directory): 132 for dirpath, dirnames, filenames in os.walk(directory):
95 for name in filenames: 133 for name in filenames:
96 filepath = str(os.path.join(dirpath, name)) 134 filepath = str(os.path.join(dirpath, name))
97 - parser = dicom.Parser()  
98 counter += 1 135 counter += 1
99 if gui: 136 if gui:
100 yield (counter,nfiles) 137 yield (counter,nfiles)
101 - if parser.SetFileName(filepath):  
102 - dcm = dicom.Dicom()  
103 - dcm.SetParser(parser)  
104 - grouper.AddFile(dcm) 138 + q.put(filepath)
105 else: 139 else:
106 dirpath, dirnames, filenames = os.walk(directory) 140 dirpath, dirnames, filenames = os.walk(directory)
107 for name in filenames: 141 for name in filenames:
108 filepath = str(os.path.join(dirpath, name)) 142 filepath = str(os.path.join(dirpath, name))
109 - parser = dicom.Parser()  
110 counter += 1 143 counter += 1
111 if gui: 144 if gui:
112 yield (counter,nfiles) 145 yield (counter,nfiles)
113 - if parser.SetFileName(filepath):  
114 - dcm = dicom.Dicom()  
115 - dcm.SetParser(parser)  
116 - grouper.AddFile(dcm) 146 + q.put(filepath)
  147 +
  148 + for t in threads:
  149 + q.put(0)
  150 +
  151 + for t in threads:
  152 + t.join()
117 153
118 yield grouper.GetPatientsGroups() 154 yield grouper.GetPatientsGroups()
119 155