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 16 # PARTICULAR. Consulte a Licenca Publica Geral GNU para obter mais
17 17 # detalhes.
18 18 #--------------------------------------------------------------------------
19   -from vtk.util.colors import yellow
20 19 import glob
21   -import os
22 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 28 import vtk
25 29 import vtkgdcm
26 30 import gdcm
27   -import thread
28 31 import wx
29 32 import wx.lib.pubsub as ps
30 33  
  34 +from vtk.util.colors import yellow
  35 +
31 36 import constants as const
32 37 import dicom
33 38 import dicom_grouper
... ... @@ -73,12 +78,38 @@ def SortFiles(filelist, dicom):
73 78  
74 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 107 def yGetDicomGroups(directory, recursive=True, gui=True):
78 108 """
79 109 Return all full paths to DICOM files inside given directory.
80 110 """
81 111 nfiles = 0
  112 + tempdir = tempfile.mkdtemp()
82 113 # Find total number of files
83 114 if recursive:
84 115 for dirpath, dirnames, filenames in os.walk(directory):
... ... @@ -89,31 +120,36 @@ def yGetDicomGroups(directory, recursive=True, gui=True):
89 120  
90 121 counter = 0
91 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 130 # Retrieve only DICOM files, splited into groups
93 131 if recursive:
94 132 for dirpath, dirnames, filenames in os.walk(directory):
95 133 for name in filenames:
96 134 filepath = str(os.path.join(dirpath, name))
97   - parser = dicom.Parser()
98 135 counter += 1
99 136 if gui:
100 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 139 else:
106 140 dirpath, dirnames, filenames = os.walk(directory)
107 141 for name in filenames:
108 142 filepath = str(os.path.join(dirpath, name))
109   - parser = dicom.Parser()
110 143 counter += 1
111 144 if gui:
112 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 154 yield grouper.GetPatientsGroups()
119 155  
... ...