Commit bdb92e42ff0bd165dfeb6ca4ef40b9e51c64044b
1 parent
f635c217
Exists in
master
and in
68 other branches
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 | ... | ... |