class EMTomoBoxer(QtGui.QMainWindow): """This class represents the EMTomoBoxer application instance. """ keypress = QtCore.pyqtSignal(QtGui.QKeyEvent) module_closed = QtCore.pyqtSignal() def __init__(self,application,options,datafile): QtGui.QWidget.__init__(self) self.initialized=False self.app=weakref.ref(application) self.options=options self.yshort=False self.apix=options.apix self.currentset=0 self.shrink=1#options.shrink self.setWindowTitle("Main Window (e2spt_boxer.py)") if options.mode=="3D": self.boxshape="circle" else: self.boxshape="rect" # Menu Bar self.mfile=self.menuBar().addMenu("File") self.mfile_open=self.mfile.addAction("Open") self.mfile_read_boxloc=self.mfile.addAction("Read Box Coord") self.mfile_save_boxloc=self.mfile.addAction("Save Box Coord") self.mfile_save_boxes_stack=self.mfile.addAction("Save Boxes as Stack") self.mfile_quit=self.mfile.addAction("Quit") self.setCentralWidget(QtGui.QWidget()) self.gbl = QtGui.QGridLayout(self.centralWidget()) # relative stretch factors self.gbl.setColumnStretch(0,1) self.gbl.setColumnStretch(1,4) self.gbl.setColumnStretch(2,0) self.gbl.setRowStretch(1,1) self.gbl.setRowStretch(0,4) # 3 orthogonal restricted projection views self.xyview = EMImage2DWidget() self.gbl.addWidget(self.xyview,0,1) self.xzview = EMImage2DWidget() self.gbl.addWidget(self.xzview,1,1) self.zyview = EMImage2DWidget() self.gbl.addWidget(self.zyview,0,0) # Select Z for xy view self.wdepth = QtGui.QSlider() self.gbl.addWidget(self.wdepth,1,2) ### Control panel area in upper left corner self.gbl2 = QtGui.QGridLayout() self.gbl.addLayout(self.gbl2,1,0) #self.wxpos = QtGui.QSlider(Qt.Horizontal) #self.gbl2.addWidget(self.wxpos,0,0) #self.wypos = QtGui.QSlider(Qt.Vertical) #self.gbl2.addWidget(self.wypos,0,3,6,1) # box size self.wboxsize=ValBox(label="Box Size:",value=0) self.gbl2.addWidget(self.wboxsize,2,0,1,2) # max or mean #self.wmaxmean=QtGui.QPushButton("MaxProj") #self.wmaxmean.setCheckable(True) #self.gbl2.addWidget(self.wmaxmean,3,0) # number slices self.wnlayers=QtGui.QSpinBox() self.wnlayers.setMinimum(1) self.wnlayers.setMaximum(256) self.wnlayers.setValue(1) self.gbl2.addWidget(self.wnlayers,3,1) # Local boxes in side view self.wlocalbox=QtGui.QCheckBox("Limit Side Boxes") self.gbl2.addWidget(self.wlocalbox,3,0) self.wlocalbox.setChecked(True) # scale factor self.wscale=ValSlider(rng=(.1,2),label="Sca:",value=1.0) self.gbl2.addWidget(self.wscale,4,0,1,2) # 2-D filters self.wfilt = ValSlider(rng=(0,150),label="Filt:",value=0.0) self.gbl2.addWidget(self.wfilt,5,0,1,2) self.curbox=-1 self.boxes=[] # array of box info, each is (x,y,z,...) self.boxesimgs=[] # z projection of each box self.xydown=self.xzdown=self.zydown=None self.firsthbclick = None # coordinate display self.wcoords=QtGui.QLabel("X: " + str(self.get_x()) + "\t\t" + "Y: " + str(self.get_y()) + "\t\t" + "Z: " + str(self.get_z())) self.gbl2.addWidget(self.wcoords, 1, 0, 1, 2) # file menu self.mfile_open.triggered[bool].connect(self.menu_file_open) self.mfile_read_boxloc.triggered[bool].connect(self.menu_file_read_boxloc) self.mfile_save_boxloc.triggered[bool].connect(self.menu_file_save_boxloc) self.mfile_save_boxes_stack.triggered[bool].connect(self.save_boxes) self.mfile_quit.triggered[bool].connect(self.menu_file_quit) # all other widgets self.wdepth.valueChanged[int].connect(self.event_depth) self.wnlayers.valueChanged[int].connect(self.event_nlayers) self.wboxsize.valueChanged.connect(self.event_boxsize) #self.wmaxmean.clicked[bool].connect(self.event_projmode) self.wscale.valueChanged.connect(self.event_scale) self.wfilt.valueChanged.connect(self.event_filter) self.wlocalbox.stateChanged[int].connect(self.event_localbox) self.xyview.mousemove.connect(self.xy_move) self.xyview.mousedown.connect(self.xy_down) self.xyview.mousedrag.connect(self.xy_drag) self.xyview.mouseup.connect(self.xy_up) self.xyview.mousewheel.connect(self.xy_wheel) self.xyview.signal_set_scale.connect(self.xy_scale) self.xyview.origin_update.connect(self.xy_origin) self.xzview.mousedown.connect(self.xz_down) self.xzview.mousedrag.connect(self.xz_drag) self.xzview.mouseup.connect(self.xz_up) self.xzview.signal_set_scale.connect(self.xz_scale) self.xzview.origin_update.connect(self.xz_origin) self.zyview.mousedown.connect(self.zy_down) self.zyview.mousedrag.connect(self.zy_drag) self.zyview.mouseup.connect(self.zy_up) self.zyview.signal_set_scale.connect(self.zy_scale) self.zyview.origin_update.connect(self.zy_origin) self.xyview.keypress.connect(self.key_press) self.datafilename=datafile self.basename=base_name(datafile) p0=datafile.find('__') if p0>0: p1=datafile.rfind('.') self.filetag=datafile[p0:p1] if self.filetag[-1]!='_': self.filetag+='_' else: self.filetag="__" data=EMData(datafile) self.set_data(data) # Boxviewer subwidget (details of a single box) self.boxviewer=EMBoxViewer() #self.app().attach_child(self.boxviewer) # Boxes Viewer (z projections of all boxes) self.boxesviewer=EMImageMXWidget() #self.app().attach_child(self.boxesviewer) self.boxesviewer.show() self.boxesviewer.set_mouse_mode("App") self.boxesviewer.setWindowTitle("Particle List") self.boxesviewer.rzonce=True self.setspanel=EMTomoSetsPanel(self) self.optionviewer=EMTomoBoxerOptions(self) self.optionviewer.add_panel(self.setspanel,"Sets") self.optionviewer.show() # Average viewer shows results of background tomographic processing # self.averageviewer=EMAverageViewer(self) #self.averageviewer.show() self.boxesviewer.mx_image_selected.connect(self.img_selected) self.jsonfile=info_name(datafile) info=js_open_dict(self.jsonfile) self.sets={} self.boxsize={} if "class_list" in info: clslst=info["class_list"] for k in sorted(clslst.keys()): if type(clslst[k])==dict: self.sets[int(k)]=str(clslst[k]["name"]) self.boxsize[int(k)]=int(clslst[k]["boxsize"]) else: self.sets[int(k)]=str(clslst[k]) self.boxsize[int(k)]=boxsize clr=QtGui.QColor self.setcolors=[clr("blue"),clr("green"),clr("red"),clr("cyan"),clr("purple"),clr("orange"), clr("yellow"),clr("hotpink"),clr("gold")] self.sets_visible={} if "boxes_3d" in info: box=info["boxes_3d"] for i,b in enumerate(box): #### X-center,Y-center,Z-center,method,[score,[class #]] bdf=[0,0,0,"manual",0.0, 0] for j,bi in enumerate(b): bdf[j]=bi if bdf[5] not in list(self.sets.keys()): clsi=int(bdf[5]) self.sets[clsi]="particles_{:02d}".format(clsi) self.boxsize[clsi]=boxsize self.boxes.append(bdf) ###### this is the new (2018-09) metadata standard.. ### now we use coordinates at full size from center of tomogram so it works for different binning and clipping ### have to make it compatible with older versions though.. if "apix_unbin" in info: self.apix_unbin=info["apix_unbin"] self.apix_cur=apix=data["apix_x"] for b in self.boxes: b[0]=b[0]/apix*self.apix_unbin+data["nx"]//2 b[1]=b[1]/apix*self.apix_unbin+data["ny"]//2 b[2]=b[2]/apix*self.apix_unbin+data["nz"]//2 for k in self.boxsize.keys(): self.boxsize[k]=self.boxsize[k]/apix*self.apix_unbin else: self.apix_unbin=-1 info.close() if len(self.sets)==0: self.new_set("particles_00") self.sets_visible[list(self.sets.keys())[0]]=0 self.currentset=sorted(self.sets.keys())[0] self.setspanel.update_sets() self.wboxsize.setValue(self.get_boxsize()) print(self.sets) for i in range(len(self.boxes)): self.update_box(i) self.update_all() self.initialized=True def set_data(self,data): self.data=data self.apix=data["apix_x"] self.datasize=(data["nx"],data["ny"],data["nz"]) self.wdepth.setRange(0,self.datasize[2]-1) self.boxes=[] self.curbox=-1 self.wdepth.setValue(old_div(self.datasize[2],2)) if self.initialized: self.update_all() def eraser_width(self): return int(self.optionviewer.eraser_radius.getValue()) def get_cube(self,x,y,z, centerslice=False, boxsz=-1): """Returns a box-sized cube at the given center location""" if boxsz<0: bs=self.get_boxsize() else: bs=boxsz if centerslice: bz=1 else: bz=bs if ((x<-bs//2) or (y<-bs//2) or (z<-bz//2) or (x>self.data["nx"]+bs//2) or (y>self.data["ny"]+bs//2) or (z>self.data["nz"]+bz//2) ): r=EMData(bs,bs,bz) else: r=self.data.get_clip(Region(x-bs//2,y-bs//2,z-bz//2,bs,bs,bz)) if self.apix!=0 : r["apix_x"]=self.apix r["apix_y"]=self.apix r["apix_z"]=self.apix #if options.normproc: #r.process_inplace(options.normproc) return r def get_slice(self,n,xyz): """Reads a slice either from a file or the preloaded memory array. xyz is the axis along which 'n' runs, 0=x (yz), 1=y (xz), 2=z (xy)""" if xyz==0: r=self.data.get_clip(Region(n,0,0,1,self.datasize[1],self.datasize[2])) r.set_size(self.datasize[1],self.datasize[2],1) elif xyz==1: r=self.data.get_clip(Region(0,n,0,self.datasize[0],1,self.datasize[2])) r.set_size(self.datasize[0],self.datasize[2],1) else: r=self.data.get_clip(Region(0,0,n,self.datasize[0],self.datasize[1],1)) if self.apix!=0 : r["apix_x"]=self.apix r["apix_y"]=self.apix r["apix_z"]=self.apix return r def event_boxsize(self): if self.get_boxsize()==int(self.wboxsize.getValue()): return self.boxsize[self.currentset]=int(self.wboxsize.getValue()) cb=self.curbox self.initialized=False for i in range(len(self.boxes)): if self.boxes[i][5]==self.currentset: self.update_box(i) self.update_box(cb) self.initialized=True self.update_all() def event_projmode(self,state): """Projection mode can be simple average (state=False) or maximum projection (state=True)""" self.update_all() def event_scale(self,newscale): self.xyview.set_scale(newscale) self.xzview.set_scale(newscale) self.zyview.set_scale(newscale) def event_depth(self): if self.initialized: self.update_xy() def event_nlayers(self): self.update_all() def event_filter(self): self.update_all() def event_localbox(self,tog): self.update_sides() def get_boxsize(self, clsid=-1): if clsid<0: return int(self.boxsize[self.currentset]) else: try: ret= int(self.boxsize[clsid]) except: print("No box size saved for {}..".format(clsid)) ret=32 return ret def nlayers(self): return int(self.wnlayers.value()) def depth(self): return int(self.wdepth.value()) def scale(self): return self.wscale.getValue() def get_x(self): return self.get_coord(0) def get_y(self): return self.get_coord(1) def get_z(self): return self.depth() def get_coord(self, coord_index): if len(self.boxes) > 1: if self.curbox: return self.boxes[self.curbox][coord_index] else: return self.boxes[-1][coord_index] else: return 0 def menu_file_open(self,tog): QtGui.QMessageBox.warning(None,"Error","Sorry, in the current version, you must provide a file to open on the command-line.") def load_box_yshort(self, boxcoords): if options.yshort: return [boxcoords[0], boxcoords[2], boxcoords[1]] else: return boxcoords def menu_file_read_boxloc(self): fsp=str(QtGui.QFileDialog.getOpenFileName(self, "Select output text file")) f=file(fsp,"r") for b in f: b2=[old_div(int(float(i)),self.shrink) for i in b.split()[:3]] bdf=[0,0,0,"manual",0.0, self.currentset] for j in range(len(b2)): bdf[j]=b2[j] self.boxes.append(bdf) self.update_box(len(self.boxes)-1) f.close() def menu_file_save_boxloc(self): shrinkf=self.shrink #jesus fsp=str(QtGui.QFileDialog.getSaveFileName(self, "Select output text file")) out=file(fsp,"w") for b in self.boxes: out.write("%d\t%d\t%d\n"%(b[0]*shrinkf,b[1]*shrinkf,b[2]*shrinkf)) out.close() def save_boxes(self, clsid=[]): if len(clsid)==0: defaultname="ptcls.hdf" else: defaultname="_".join([self.sets[i] for i in clsid])+".hdf" name,ok=QtGui.QInputDialog.getText( self, "Save particles", "Filename suffix:", text=defaultname) if not ok: return name=self.filetag+str(name) if name[-4:].lower()!=".hdf" : name+=".hdf" if self.options.mode=="3D": dr="particles3d" is2d=False else: dr="particles" is2d=True if not os.path.isdir(dr): os.mkdir(dr) fsp=os.path.join(dr,self.basename)+name print("Saving {} particles to {}".format(self.options.mode, fsp)) if os.path.isfile(fsp): print("{} exist. Overwritting...".format(fsp)) os.remove(fsp) progress = QtGui.QProgressDialog("Saving", "Abort", 0, len(self.boxes),None) boxsz=-1 for i,b in enumerate(self.boxes): if len(clsid)>0: if int(b[5]) not in clsid: continue #img=self.get_cube(b[0],b[1],b[2]) bs=self.get_boxsize(b[5]) if boxsz<0: boxsz=bs else: if boxsz!=bs: print("Inconsistant box size in the particles to save.. Using {:d}..".format(boxsz)) bs=boxsz sz=[s//2 for s in self.datasize] img=self.get_cube(b[0], b[1], b[2], centerslice=is2d, boxsz=bs) if is2d==False: img.process_inplace('normalize') img["ptcl_source_image"]=self.datafilename img["ptcl_source_coord"]=(b[0]-sz[0], b[1]-sz[1], b[2]-sz[2]) if is2d==False: #### do not invert contrast for 2D images img.mult(-1) img.write_image(fsp,-1) progress.setValue(i+1) if progress.wasCanceled(): break def menu_file_quit(self): self.close() def transform_coords(self, point, xform): xvec = xform.get_matrix() return [xvec[0]*point[0] + xvec[4]*point[1] + xvec[8]*point[2] + xvec[3], xvec[1]*point[0] + xvec[5]*point[1] + xvec[9]*point[2] + xvec[7], xvec[2]*point[0] + xvec[6]*point[1] + xvec[10]*point[2] + xvec[11]] def get_averager(self): """returns an averager of the appropriate type for generating projection views""" #if self.wmaxmean.isChecked() : return Averagers.get("minmax",{"max":1}) return Averagers.get("mean") def update_sides(self): """updates xz and yz views due to a new center location""" #print "\n\n\n\n\nIn update sides, self.datafile is", self.datafile #print "\n\n\n\n" if self.data==None: return if self.curbox==-1 : x=self.datasize[0]//2 y=self.datasize[1]//2 z=0 else: x,y,z=self.boxes[self.curbox][:3] self.cury=y self.curx=x # update shape display if self.wlocalbox.isChecked(): xzs=self.xzview.get_shapes() for i in range(len(self.boxes)): bs=self.get_boxsize(self.boxes[i][5]) if self.boxes[i][1]<self.cury+old_div(bs,2) and self.boxes[i][1]>self.cury-old_div(bs,2) and self.boxes[i][5] in self.sets_visible: xzs[i][0]=self.boxshape else: xzs[i][0]="hidden" zys=self.zyview.get_shapes() for i in range(len(self.boxes)): bs=self.get_boxsize(self.boxes[i][5]) if self.boxes[i][0]<self.curx+old_div(bs,2) and self.boxes[i][0]>self.curx-old_div(bs,2) and self.boxes[i][5] in self.sets_visible: zys[i][0]=self.boxshape else: zys[i][0]="hidden" else : xzs=self.xzview.get_shapes() zys=self.zyview.get_shapes() for i in range(len(self.boxes)): bs=self.get_boxsize(self.boxes[i][5]) if self.boxes[i][5] in self.sets_visible: xzs[i][0]=self.boxshape zys[i][0]=self.boxshape else: xzs[i][0]="hidden" zys[i][0]="hidden" self.xzview.shapechange=1 self.zyview.shapechange=1 # yz avgr=self.get_averager() for x in range(x-(self.nlayers()//2),x+((self.nlayers()+1)//2)): slc=self.get_slice(x,0) avgr.add_image(slc) av=avgr.finish() if not self.yshort: av.process_inplace("xform.transpose") if self.wfilt.getValue()!=0.0: av.process_inplace("filter.lowpass.gauss",{"cutoff_freq":old_div(1.0,self.wfilt.getValue()),"apix":self.apix}) self.zyview.set_data(av) # xz avgr=self.get_averager() for y in range(y-old_div(self.nlayers(),2),y+old_div((self.nlayers()+1),2)): slc=self.get_slice(y,1) avgr.add_image(slc) av=avgr.finish() if self.wfilt.getValue()!=0.0: av.process_inplace("filter.lowpass.gauss",{"cutoff_freq":old_div(1.0,self.wfilt.getValue()),"apix":self.apix}) self.xzview.set_data(av) def update_xy(self): """updates xy view due to a new slice range""" #print "\n\n\n\n\nIn update_xy, self.datafile is", self.datafile #print "\n\n\n\n" if self.data==None: return # Boxes should also be limited by default in the XY view if len(self.boxes) > 0: zc=self.wdepth.value() #print "The current depth is", self.wdepth.value() xys=self.xyview.get_shapes() for i in range(len(self.boxes)): bs=self.get_boxsize(self.boxes[i][5]) zdist=abs(self.boxes[i][2] - zc) if self.options.mode=="3D": zthr=bs/2 xys[i][6]=bs//2-zdist else: zthr=1 if zdist < zthr and self.boxes[i][5] in self.sets_visible: xys[i][0]=self.boxshape else : xys[i][0]="hidden" self.xyview.shapechange=1 #if self.wmaxmean.isChecked(): #avgr=Averagers.get("minmax",{"max":1}) #else: avgr=Averagers.get("mean") slc=EMData() for z in range(self.wdepth.value()-self.nlayers()//2,self.wdepth.value()+(self.nlayers()+1)//2): slc=self.get_slice(z,2) avgr.add_image(slc) av=avgr.finish() #print "\n\nIn update xy, av and type are", av, type(av) if self.wfilt.getValue()!=0.0: av.process_inplace("filter.lowpass.gauss",{"cutoff_freq":old_div(1.0,self.wfilt.getValue()),"apix":self.apix}) if self.initialized: self.xyview.set_data(av, keepcontrast=True) else: self.xyview.set_data(av) def update_all(self): """redisplay of all widgets""" #print "\n\n\n\n\nIn update all, self.datafile is", self.datafile #print "\n\n\n\n" if self.data==None: return self.update_xy() self.update_sides() self.update_boximgs() def update_coords(self): self.wcoords.setText("X: " + str(self.get_x()) + "\t\t" + "Y: " + str(self.get_y()) + "\t\t" + "Z: " + str(self.get_z())) def inside_box(self,n,x=-1,y=-1,z=-1): """Checks to see if a point in image coordinates is inside box number n. If any value is negative, it will not be checked.""" box=self.boxes[n] if box[5] not in self.sets_visible: return False bs=self.get_boxsize(box[5])/2 if self.options.mode=="3D": rr=(x>=0)*((box[0]-x)**2) + (y>=0)*((box[1]-y) **2) + (z>=0)*((box[2]-z)**2) else: rr=(x>=0)*((box[0]-x)**2) + (y>=0)*((box[1]-y) **2) + (z>=0)*(box[2]!=z)*(1e3*bs**2) return rr<=bs**2 def do_deletion(self, delids): kpids=[i for i,b in enumerate(self.boxes) if i not in delids] self.boxes=[self.boxes[i] for i in kpids] self.boxesimgs=[self.boxesimgs[i] for i in kpids] self.xyview.shapes={i:self.xyview.shapes[k] for i,k in enumerate(kpids)} self.xzview.shapes={i:self.xzview.shapes[k] for i,k in enumerate(kpids)} self.zyview.shapes={i:self.zyview.shapes[k] for i,k in enumerate(kpids)} #print self.boxes, self.xyview.get_shapes() self.curbox=-1 self.update_all() def del_box(self,n): """Delete an existing box by replacing the deleted box with the last box. A bit funny, but otherwise update after deletion is REALLY slow.""" # print "del ",n if n<0 or n>=len(self.boxes): return if self.boxviewer.get_data(): self.boxviewer.set_data(None) self.curbox=-1 self.do_deletion([n]) def update_box(self,n,quiet=False): """After adjusting a box, call this""" # print "upd ",n,quiet try: box=self.boxes[n] except IndexError: return bs2=self.get_boxsize(box[5])//2 color=self.setcolors[box[5]%len(self.setcolors)].getRgbF() if self.options.mode=="3D": self.xyview.add_shape(n,EMShape(["circle",color[0],color[1],color[2],box[0],box[1],bs2,2])) self.xzview.add_shape(n,EMShape(["circle",color[0],color[1],color[2],box[0],box[2],bs2,2])) self.zyview.add_shape(n,EMShape(("circle",color[0],color[1],color[2],box[2],box[1],bs2,2))) else: self.xyview.add_shape(n,EMShape(["rect",color[0],color[1],color[2], box[0]-bs2,box[1]-bs2,box[0]+bs2,box[1]+bs2,2])) self.xzview.add_shape(n,EMShape(["rect",color[0],color[1],color[2], box[0]-bs2,box[2]-1,box[0]+bs2,box[2]+1,2])) self.zyview.add_shape(n,EMShape(["rect",color[0],color[1],color[2], box[2]-1,box[1]-bs2,box[2]+1,box[1]+bs2,2])) if self.depth()!=box[2]: self.wdepth.setValue(box[2]) else: self.xyview.update() if self.initialized: self.update_sides() # For speed, we turn off updates while dragging a box around. Quiet is set until the mouse-up if not quiet: # Get the cube from the original data (normalized) proj=self.get_cube(box[0], box[1], box[2], centerslice=True, boxsz=self.get_boxsize(box[5])) proj.process_inplace("normalize") for i in range(len(self.boxesimgs),n+1): self.boxesimgs.append(None) self.boxesimgs[n]=proj mm=[m for im,m in enumerate(self.boxesimgs) if self.boxes[im][5] in self.sets_visible] if self.initialized: self.SaveJson() if self.initialized: self.update_boximgs() if n!=self.curbox: self.boxesviewer.set_selected((n,),True) self.curbox=n self.update_coords() def update_boximgs(self): self.boxids=[im for im,m in enumerate(self.boxesimgs) if self.boxes[im][5] in self.sets_visible] self.boxesviewer.set_data([self.boxesimgs[i] for i in self.boxids]) self.boxesviewer.update() return def img_selected(self,event,lc): #print "sel",lc[0] lci=self.boxids[lc[0]] if event.modifiers()&Qt.ShiftModifier: self.del_box(lci) else: self.update_box(lci) if self.curbox>=0 : box=self.boxes[self.curbox] self.xyview.scroll_to(box[0],box[1]) self.xzview.scroll_to(None,box[2]) self.zyview.scroll_to(box[2],None) self.currentset=box[5] self.setspanel.initialized=False self.setspanel.update_sets() def del_region_xy(self, x=-1, y=-1, z=-1, rad=-1): if rad<0: rad=self.eraser_width() delids=[] for i,b in enumerate(self.boxes): if b[5] not in self.sets_visible: continue if (x>=0)*(b[0]-x)**2 + (y>=0)*(b[1]-y)**2 +(z>=0)*(b[2]-z)**2 < rad**2: delids.append(i) self.do_deletion(delids) def xy_down(self,event): x,y=self.xyview.scr_to_img((event.x(),event.y())) x,y=int(x),int(y) z=int(self.get_z()) self.xydown=None if x<0 or y<0 : return # no clicking outside the image (on 2 sides) if self.optionviewer.erasercheckbox.isChecked(): self.del_region_xy(x,y) return for i in range(len(self.boxes)): if self.inside_box(i,x,y,z): if event.modifiers()&Qt.ShiftModifier: self.del_box(i) self.firsthbclick = None else: self.xydown=(i,x,y,self.boxes[i][0],self.boxes[i][1]) self.update_box(i) break else: # if x>self.get_boxsize()/2 and x<self.datasize[0]-self.get_boxsize()/2 and y>self.get_boxsize()/2 and y<self.datasize[1]-self.get_boxsize()/2 and self.depth()>self.get_boxsize()/2 and self.depth()<self.datasize[2]-self.get_boxsize()/2 : if not event.modifiers()&Qt.ShiftModifier: self.boxes.append(([x,y,self.depth(), 'manual', 0.0, self.currentset])) self.xydown=(len(self.boxes)-1,x,y,x,y) # box #, x down, y down, x box at down, y box at down self.update_box(self.xydown[0]) if self.curbox>=0: box=self.boxes[self.curbox] self.xzview.scroll_to(None,box[2]) self.zyview.scroll_to(box[2],None) def xy_drag(self,event): x,y=self.xyview.scr_to_img((event.x(),event.y())) x,y=int(x),int(y) if self.optionviewer.erasercheckbox.isChecked(): self.del_region_xy(x,y) self.xyview.eraser_shape=EMShape(["circle",1,1,1,x,y,self.eraser_width(),2]) self.xyview.shapechange=1 self.xyview.update() return if self.xydown==None : return dx=x-self.xydown[1] dy=y-self.xydown[2] self.boxes[self.xydown[0]][0]=dx+self.xydown[3] self.boxes[self.xydown[0]][1]=dy+self.xydown[4] self.update_box(self.curbox,True) def xy_up (self,event): if self.xydown!=None: self.update_box(self.curbox) self.xydown=None def xy_wheel (self,event): if event.delta() > 0: #self.wdepth.setValue(self.wdepth.value()+4) self.wdepth.setValue(self.wdepth.value()+1) #jesus elif event.delta() < 0: #self.wdepth.setValue(self.wdepth.value()-4) self.wdepth.setValue(self.wdepth.value()-1) #jesus def xy_scale(self,news): "xy image view has been rescaled" self.wscale.setValue(news) #self.xzview.set_scale(news,True) #self.zyview.set_scale(news,True) def xy_origin(self,newor): "xy origin change" xzo=self.xzview.get_origin() self.xzview.set_origin(newor[0],xzo[1],True) zyo=self.zyview.get_origin() self.zyview.set_origin(zyo[0],newor[1],True) def xy_move(self,event): if self.optionviewer.erasercheckbox.isChecked(): x,y=self.xyview.scr_to_img((event.x(),event.y())) #print x,y self.xyview.eraser_shape=EMShape(["circle",1,1,1,x,y,self.eraser_width(),2]) self.xyview.shapechange=1 self.xyview.update() else: self.xyview.eraser_shape=None def xz_down(self,event): x,z=self.xzview.scr_to_img((event.x(),event.y())) x,z=int(x),int(z) y=int(self.get_y()) self.xzdown=None if x<0 or z<0 : return # no clicking outside the image (on 2 sides) if self.optionviewer.erasercheckbox.isChecked(): return for i in range(len(self.boxes)): if (not self.wlocalbox.isChecked() and self.inside_box(i,x,y,z)) or self.inside_box(i,x,self.cury,z) : if event.modifiers()&Qt.ShiftModifier: self.del_box(i) self.firsthbclick = None else : self.xzdown=(i,x,z,self.boxes[i][0],self.boxes[i][2]) self.update_box(i) break else: if not event.modifiers()&Qt.ShiftModifier: self.boxes.append(([x,self.cury,z, 'manual', 0.0, self.currentset])) self.xzdown=(len(self.boxes)-1,x,z,x,z) # box #, x down, y down, x box at down, y box at down self.update_box(self.xzdown[0]) if self.curbox>=0 : box=self.boxes[self.curbox] self.xyview.scroll_to(None,box[1]) self.zyview.scroll_to(box[2],None) def xz_drag(self,event): if self.xzdown==None : return x,z=self.xzview.scr_to_img((event.x(),event.y())) x,z=int(x),int(z) dx=x-self.xzdown[1] dz=z-self.xzdown[2] self.boxes[self.xzdown[0]][0]=dx+self.xzdown[3] self.boxes[self.xzdown[0]][2]=dz+self.xzdown[4] self.update_box(self.curbox,True) def xz_up (self,event): if self.xzdown!=None: self.update_box(self.curbox) self.xzdown=None def xz_scale(self,news): "xy image view has been rescaled" self.wscale.setValue(news) #self.xyview.set_scale(news,True) #self.zyview.set_scale(news,True) def xz_origin(self,newor): "xy origin change" xyo=self.xyview.get_origin() self.xyview.set_origin(newor[0],xyo[1],True) #zyo=self.zyview.get_origin() #self.zyview.set_origin(zyo[0],newor[1],True) def zy_down(self,event): z,y=self.zyview.scr_to_img((event.x(),event.y())) z,y=int(z),int(y) x=int(self.get_x()) self.xydown=None if z<0 or y<0 : return # no clicking outside the image (on 2 sides) for i in range(len(self.boxes)): if (not self.wlocalbox.isChecked() and self.inside_box(i,x,y,z)) or self.inside_box(i,self.curx,y,z): if event.modifiers()&Qt.ShiftModifier: self.del_box(i) self.firsthbclick = None else : self.zydown=(i,z,y,self.boxes[i][2],self.boxes[i][1]) self.update_box(i) break else: if not event.modifiers()&Qt.ShiftModifier: ########### self.boxes.append(([self.curx,y,z, 'manual', 0.0, self.currentset])) self.zydown=(len(self.boxes)-1,z,y,z,y) # box #, x down, y down, x box at down, y box at down self.update_box(self.zydown[0]) if self.curbox>=0 : box=self.boxes[self.curbox] self.xyview.scroll_to(box[0],None) self.xzview.scroll_to(None,box[2]) def zy_drag(self,event): if self.zydown==None : return z,y=self.zyview.scr_to_img((event.x(),event.y())) z,y=int(z),int(y) dz=z-self.zydown[1] dy=y-self.zydown[2] self.boxes[self.zydown[0]][2]=dz+self.zydown[3] self.boxes[self.zydown[0]][1]=dy+self.zydown[4] self.update_box(self.curbox,True) def zy_up (self,event): if self.zydown!=None: self.update_box(self.curbox) self.zydown=None def zy_scale(self,news): "xy image view has been rescaled" self.wscale.setValue(news) #self.xyview.set_scale(news,True) #self.xzview.set_scale(news,True) def zy_origin(self,newor): "xy origin change" xyo=self.xyview.get_origin() self.xyview.set_origin(xyo[0],newor[1],True) #xzo=self.xzview.get_origin() #self.xzview.set_origin(xzo[0],newor[1],True) def set_current_set(self, name): #print "set current", name name=parse_setname(name) self.currentset=name self.wboxsize.setValue(self.get_boxsize()) self.update_all() return def hide_set(self, name): name=parse_setname(name) if name in self.sets_visible: self.sets_visible.pop(name) if self.initialized: self.update_all() self.update_boximgs() return def show_set(self, name): name=parse_setname(name) self.sets_visible[name]=0 #self.currentset=name self.wboxsize.setValue(self.get_boxsize()) if self.initialized: self.update_all() self.update_boximgs() return def delete_set(self, name): name=parse_setname(name) ## idx to keep delids=[i for i,b in enumerate(self.boxes) if b[5]==int(name)] self.do_deletion(delids) if name in self.sets_visible: self.sets_visible.pop(name) if name in self.sets: self.sets.pop(name) if name in self.boxsize: self.boxsize.pop(name) self.curbox=-1 self.update_all() return def rename_set(self, oldname, newname): name=parse_setname(oldname) if name in self.sets: self.sets[name]=newname return def new_set(self, name): for i in range(len(self.sets)+1): if i not in self.sets: break self.sets[i]=name self.sets_visible[i]=0 if self.options.mode=="3D": self.boxsize[i]=32 else: self.boxsize[i]=64 return def save_set(self): self.save_boxes(list(self.sets_visible.keys())) return def key_press(self,event): if event.key() == 96: self.wdepth.setValue(self.wdepth.value()+1) elif event.key() == 49: self.wdepth.setValue(self.wdepth.value()-1) else: self.keypress.emit(event) def SaveJson(self): info=js_open_dict(self.jsonfile) sx,sy,sz=(self.data["nx"]//2,self.data["ny"]//2,self.data["nz"]//2) if "apix_unbin" in info: bxs=[] for b0 in self.boxes: b=[ (b0[0]-sx)*self.apix_cur/self.apix_unbin, (b0[1]-sy)*self.apix_cur/self.apix_unbin, (b0[2]-sz)*self.apix_cur/self.apix_unbin, b0[3], b0[4], b0[5] ] bxs.append(b) bxsz={} for k in self.boxsize.keys(): bxsz[k]=self.boxsize[k]*self.apix_cur/self.apix_unbin else: bxs=self.boxes bxsz=self.boxsize info["boxes_3d"]=bxs clslst={} for key in list(self.sets.keys()): clslst[int(key)]={ "name":self.sets[key], "boxsize":int(bxsz[key]), } info["class_list"]=clslst info.close() def closeEvent(self,event): print("Exiting") self.SaveJson() self.boxviewer.close() self.boxesviewer.close() self.optionviewer.close() self.xyview.close() self.xzview.close() self.zyview.close() self.module_closed.emit() # this signal is important when e2ctf is being used by a program running its own event loop
class EMTomoBoxer(QtWidgets.QMainWindow): """This class represents the EMTomoBoxer application instance. """ keypress = QtCore.pyqtSignal(QtGui.QKeyEvent) module_closed = QtCore.pyqtSignal() def __init__(self,application,options,datafile): QtWidgets.QWidget.__init__(self) self.initialized=False self.app=weakref.ref(application) self.options=options self.apix=options.apix self.currentset=0 self.shrink=1#options.shrink self.setWindowTitle("Main Window (e2spt_boxer.py)") if options.mode=="3D": self.boxshape="circle" else: self.boxshape="rect" self.globalxf=Transform() # Menu Bar self.mfile=self.menuBar().addMenu("File") #self.mfile_open=self.mfile.addAction("Open") self.mfile_read_boxloc=self.mfile.addAction("Read Box Coord") self.mfile_save_boxloc=self.mfile.addAction("Save Box Coord") self.mfile_save_boxpdb=self.mfile.addAction("Save Coord as PDB") self.mfile_save_boxes_stack=self.mfile.addAction("Save Boxes as Stack") #self.mfile_quit=self.mfile.addAction("Quit") self.setCentralWidget(QtWidgets.QWidget()) self.gbl = QtWidgets.QGridLayout(self.centralWidget()) # relative stretch factors self.gbl.setColumnMinimumWidth(0,200) self.gbl.setRowMinimumHeight(0,200) self.gbl.setColumnStretch(0,0) self.gbl.setColumnStretch(1,100) self.gbl.setColumnStretch(2,0) self.gbl.setRowStretch(1,0) self.gbl.setRowStretch(0,100) # 3 orthogonal restricted projection views self.xyview = EMImage2DWidget(sizehint=(1024,1024)) self.gbl.addWidget(self.xyview,0,1) self.xzview = EMImage2DWidget(sizehint=(1024,256)) self.gbl.addWidget(self.xzview,1,1) self.zyview = EMImage2DWidget(sizehint=(256,1024)) self.gbl.addWidget(self.zyview,0,0) # Select Z for xy view self.wdepth = QtWidgets.QSlider() self.gbl.addWidget(self.wdepth,1,2) ### Control panel area in upper left corner self.gbl2 = QtWidgets.QGridLayout() self.gbl.addLayout(self.gbl2,1,0) #self.wxpos = QtWidgets.QSlider(Qt.Horizontal) #self.gbl2.addWidget(self.wxpos,0,0) #self.wypos = QtWidgets.QSlider(Qt.Vertical) #self.gbl2.addWidget(self.wypos,0,3,6,1) # box size self.wboxsize=ValBox(label="Box Size:",value=0) self.gbl2.addWidget(self.wboxsize,2,0) # max or mean #self.wmaxmean=QtWidgets.QPushButton("MaxProj") #self.wmaxmean.setCheckable(True) #self.gbl2.addWidget(self.wmaxmean,3,0) # number slices label0=QtWidgets.QLabel("Thickness") self.gbl2.addWidget(label0,3,0) self.wnlayers=QtWidgets.QSpinBox() self.wnlayers.setMinimum(1) self.wnlayers.setMaximum(256) self.wnlayers.setValue(1) self.gbl2.addWidget(self.wnlayers,3,1) # Local boxes in side view self.wlocalbox=QtWidgets.QCheckBox("Limit Side Boxes") self.gbl2.addWidget(self.wlocalbox,4,0) self.wlocalbox.setChecked(True) self.button_flat = QtWidgets.QPushButton("Flatten") self.gbl2.addWidget(self.button_flat,5,0) self.button_reset = QtWidgets.QPushButton("Reset") self.gbl2.addWidget(self.button_reset,5,1) ## scale factor #self.wscale=ValSlider(rng=(.1,2),label="Sca:",value=1.0) #self.gbl2.addWidget(self.wscale,4,0,1,2) # 2-D filters self.wfilt = ValSlider(rng=(0,150),label="Filt",value=0.0) self.gbl2.addWidget(self.wfilt,6,0,1,2) self.curbox=-1 self.boxes=[] # array of box info, each is (x,y,z,...) self.boxesimgs=[] # z projection of each box self.dragging=-1 ##coordinate display self.wcoords=QtWidgets.QLabel("") self.gbl2.addWidget(self.wcoords, 1, 0, 1, 2) self.button_flat.clicked[bool].connect(self.flatten_tomo) self.button_reset.clicked[bool].connect(self.reset_flatten_tomo) # file menu #self.mfile_open.triggered[bool].connect(self.menu_file_open) self.mfile_read_boxloc.triggered[bool].connect(self.menu_file_read_boxloc) self.mfile_save_boxloc.triggered[bool].connect(self.menu_file_save_boxloc) self.mfile_save_boxpdb.triggered[bool].connect(self.menu_file_save_boxpdb) self.mfile_save_boxes_stack.triggered[bool].connect(self.save_boxes) #self.mfile_quit.triggered[bool].connect(self.menu_file_quit) # all other widgets self.wdepth.valueChanged[int].connect(self.event_depth) self.wnlayers.valueChanged[int].connect(self.event_nlayers) self.wboxsize.valueChanged.connect(self.event_boxsize) #self.wmaxmean.clicked[bool].connect(self.event_projmode) #self.wscale.valueChanged.connect(self.event_scale) self.wfilt.valueChanged.connect(self.event_filter) self.wlocalbox.stateChanged[int].connect(self.event_localbox) self.xyview.mousemove.connect(self.xy_move) self.xyview.mousedown.connect(self.xy_down) self.xyview.mousedrag.connect(self.xy_drag) self.xyview.mouseup.connect(self.mouse_up) self.xyview.mousewheel.connect(self.xy_wheel) self.xyview.signal_set_scale.connect(self.event_scale) self.xyview.origin_update.connect(self.xy_origin) self.xzview.mousedown.connect(self.xz_down) self.xzview.mousedrag.connect(self.xz_drag) self.xzview.mouseup.connect(self.mouse_up) self.xzview.mousewheel.connect(self.xz_wheel) self.xzview.signal_set_scale.connect(self.event_scale) self.xzview.origin_update.connect(self.xz_origin) self.xzview.mousemove.connect(self.xz_move) self.zyview.mousedown.connect(self.zy_down) self.zyview.mousedrag.connect(self.zy_drag) self.zyview.mouseup.connect(self.mouse_up) self.zyview.mousewheel.connect(self.zy_wheel) self.zyview.signal_set_scale.connect(self.event_scale) self.zyview.origin_update.connect(self.zy_origin) self.zyview.mousemove.connect(self.zy_move) self.xyview.keypress.connect(self.key_press) self.datafilename=datafile self.basename=base_name(datafile) p0=datafile.find('__') if p0>0: p1=datafile.rfind('.') self.filetag=datafile[p0:p1] if self.filetag[-1]!='_': self.filetag+='_' else: self.filetag="__" data=EMData(datafile) self.set_data(data) # Boxviewer subwidget (details of a single box) #self.boxviewer=EMBoxViewer() #self.app().attach_child(self.boxviewer) # Boxes Viewer (z projections of all boxes) self.boxesviewer=EMImageMXWidget() #self.app().attach_child(self.boxesviewer) self.boxesviewer.show() self.boxesviewer.set_mouse_mode("App") self.boxesviewer.setWindowTitle("Particle List") self.boxesviewer.rzonce=True self.setspanel=EMTomoSetsPanel(self) self.optionviewer=EMTomoBoxerOptions(self) self.optionviewer.add_panel(self.setspanel,"Sets") self.optionviewer.show() self.boxesviewer.mx_image_selected.connect(self.img_selected) ################## #### deal with metadata in the _info.json file... self.jsonfile=info_name(datafile) info=js_open_dict(self.jsonfile) #### read particle classes self.sets={} self.boxsize={} if "class_list" in info: clslst=info["class_list"] for k in sorted(clslst.keys()): if type(clslst[k])==dict: self.sets[int(k)]=str(clslst[k]["name"]) self.boxsize[int(k)]=int(clslst[k]["boxsize"]) else: self.sets[int(k)]=str(clslst[k]) self.boxsize[int(k)]=64 clr=QtGui.QColor self.setcolors=[QtGui.QBrush(clr("blue")),QtGui.QBrush(clr("green")),QtGui.QBrush(clr("red")),QtGui.QBrush(clr("cyan")),QtGui.QBrush(clr("purple")),QtGui.QBrush(clr("orange")), QtGui.QBrush(clr("yellow")),QtGui.QBrush(clr("hotpink")),QtGui.QBrush(clr("gold"))] self.sets_visible={} #### read boxes if "boxes_3d" in info: box=info["boxes_3d"] for i,b in enumerate(box): #### X-center,Y-center,Z-center,method,[score,[class #]] bdf=[0,0,0,"manual",0.0, 0] for j,bi in enumerate(b): bdf[j]=bi if bdf[5] not in list(self.sets.keys()): clsi=int(bdf[5]) self.sets[clsi]="particles_{:02d}".format(clsi) self.boxsize[clsi]=64 self.boxes.append(bdf) ###### this is the new (2018-09) metadata standard.. ### now we use coordinates at full size from center of tomogram so it works for different binning and clipping ### have to make it compatible with older versions though.. if "apix_unbin" in info: self.apix_unbin=info["apix_unbin"] self.apix_cur=apix=data["apix_x"] for b in self.boxes: b[0]=b[0]/apix*self.apix_unbin+data["nx"]//2 b[1]=b[1]/apix*self.apix_unbin+data["ny"]//2 b[2]=b[2]/apix*self.apix_unbin+data["nz"]//2 for k in self.boxsize.keys(): self.boxsize[k]=int(np.round(self.boxsize[k]*self.apix_unbin/apix)) else: self.apix_unbin=-1 info.close() E2loadappwin("e2sptboxer","main",self) E2loadappwin("e2sptboxer","boxes",self.boxesviewer.qt_parent) E2loadappwin("e2sptboxer","option",self.optionviewer) #### particle classes if len(self.sets)==0: self.new_set("particles_00") self.sets_visible[list(self.sets.keys())[0]]=0 self.currentset=sorted(self.sets.keys())[0] self.setspanel.update_sets() self.wboxsize.setValue(self.get_boxsize()) #print(self.sets) for i in range(len(self.boxes)): self.update_box(i) self.update_all() self.initialized=True def set_data(self,data): self.data=data self.apix=data["apix_x"] self.datasize=(data["nx"],data["ny"],data["nz"]) self.x_loc, self.y_loc, self.z_loc=data["nx"]//2,data["ny"]//2,data["nz"]//2 self.gbl.setRowMinimumHeight(1,max(250,data["nz"])) self.gbl.setColumnMinimumWidth(0,max(250,data["nz"])) print(data["nx"],data["ny"],data["nz"]) self.wdepth.setRange(0,data["nz"]-1) self.wdepth.setValue(data["nz"]//2) self.boxes=[] self.curbox=-1 if self.initialized: self.update_all() def eraser_width(self): return int(self.optionviewer.eraser_radius.getValue()) def get_cube(self,x,y,z, centerslice=False, boxsz=-1): """Returns a box-sized cube at the given center location""" if boxsz<0: bs=self.get_boxsize() else: bs=boxsz if centerslice: bz=1 else: bz=bs if ((x<-bs//2) or (y<-bs//2) or (z<-bz//2) or (x>self.data["nx"]+bs//2) or (y>self.data["ny"]+bs//2) or (z>self.data["nz"]+bz//2) ): r=EMData(bs,bs,bz) else: r=self.data.get_clip(Region(x-bs//2,y-bs//2,z-bz//2,bs,bs,bz)) if self.apix!=0 : r["apix_x"]=r["apix_y"]=r["apix_z"]=self.apix return r def get_slice(self,idx,thk=1,axis="z"): if self.globalxf.is_identity(): data=self.data else: data=self.dataxf t=int(thk-1) idx=int(idx) r=data.process("misc.directional_sum",{"axis":axis,"first":idx-t,"last":idx+t}) r.div(t*2+1) if self.apix!=0 : r["apix_x"]=r["apix_y"]=r["apix_z"]=self.apix return r def event_boxsize(self): if self.get_boxsize()==int(self.wboxsize.getValue()): return self.boxsize[self.currentset]=int(self.wboxsize.getValue()) #cb=self.curbox self.initialized=False for i in range(len(self.boxes)): if self.boxes[i][5]==self.currentset: self.update_box(i) #self.update_box(cb) self.initialized=True self.update_all() #def event_projmode(self,state): #"""Projection mode can be simple average (state=False) or maximum projection (state=True)""" #self.update_all() def event_scale(self,newscale): self.xyview.set_scale(newscale) self.xzview.set_scale(newscale) self.zyview.set_scale(newscale) def event_depth(self): if self.z_loc!=self.wdepth.value(): self.z_loc=self.wdepth.value() if self.initialized: self.update_sliceview() def event_nlayers(self): self.update_all() def event_filter(self): self.update_all() def event_localbox(self,tog): self.update_sliceview(['x','y']) def get_boxsize(self, clsid=-1): if clsid<0: return int(self.boxsize[self.currentset]) else: try: ret= int(self.boxsize[clsid]) except: print("No box size saved for {}..".format(clsid)) ret=32 return ret def nlayers(self): return int(self.wnlayers.value()) def menu_file_read_boxloc(self): fsp=str(QtWidgets.QFileDialog.getOpenFileName(self, "Select output text file")[0]) if not os.path.isfile(fsp): print("file does not exist") return f=open(fsp,"r") for b in f: b2=[old_div(int(float(i)),self.shrink) for i in b.split()[:3]] bdf=[0,0,0,"manual",0.0, self.currentset] for j in range(len(b2)): bdf[j]=b2[j] self.boxes.append(bdf) self.update_box(len(self.boxes)-1) f.close() def menu_file_save_boxloc(self): shrinkf=self.shrink #jesus fsp=str(QtWidgets.QFileDialog.getSaveFileName(self, "Select output text file")[0]) if len(fsp)==0: return out=open(fsp,"w") for b in self.boxes: out.write("%d\t%d\t%d\n"%(b[0]*shrinkf,b[1]*shrinkf,b[2]*shrinkf)) out.close() def menu_file_save_boxpdb(self): fsp=str(QtWidgets.QFileDialog.getSaveFileName(self, "Select output PDB file", filter="PDB (*.pdb)")[0]) if len(fsp)==0: return if fsp[-4:].lower()!=".pdb" : fsp+=".pdb" clsid=list(self.sets_visible.keys()) if len(clsid)==0: print("No visible particles to save") return bxs=np.array([[b[0], b[1], b[2]] for b in self.boxes if int(b[5]) in clsid])/10 numpy2pdb(bxs, fsp) print("PDB saved to {}. Use voxel size 0.1".format(fsp)) def save_boxes(self, clsid=[]): if len(clsid)==0: defaultname="ptcls.hdf" else: defaultname="_".join([self.sets[i] for i in clsid])+".hdf" name,ok=QtWidgets.QInputDialog.getText( self, "Save particles", "Filename suffix:", text=defaultname) if not ok: return name=self.filetag+str(name) if name[-4:].lower()!=".hdf" : name+=".hdf" if self.options.mode=="3D": dr="particles3d" is2d=False else: dr="particles" is2d=True if not os.path.isdir(dr): os.mkdir(dr) fsp=os.path.join(dr,self.basename)+name print("Saving {} particles to {}".format(self.options.mode, fsp)) if os.path.isfile(fsp): print("{} exist. Overwritting...".format(fsp)) os.remove(fsp) progress = QtWidgets.QProgressDialog("Saving", "Abort", 0, len(self.boxes),None) boxsz=-1 for i,b in enumerate(self.boxes): if len(clsid)>0: if int(b[5]) not in clsid: continue #img=self.get_cube(b[0],b[1],b[2]) bs=self.get_boxsize(b[5]) if boxsz<0: boxsz=bs else: if boxsz!=bs: print("Inconsistant box size in the particles to save.. Using {:d}..".format(boxsz)) bs=boxsz sz=[s//2 for s in self.datasize] img=self.get_cube(b[0], b[1], b[2], centerslice=is2d, boxsz=bs) if is2d==False: img.process_inplace('normalize') img["ptcl_source_image"]=self.datafilename img["ptcl_source_coord"]=(b[0]-sz[0], b[1]-sz[1], b[2]-sz[2]) if is2d==False: #### do not invert contrast for 2D images img.mult(-1) img.write_image(fsp,-1) progress.setValue(i+1) if progress.wasCanceled(): break def update_sliceview(self, axis=['x','y','z']): boxes=self.get_rotated_boxes() allside=(not self.wlocalbox.isChecked()) pms={'z':[2, self.xyview, self.z_loc], 'y':[1, self.xzview, self.y_loc], 'x':[0, self.zyview, self.x_loc]} if self.boxshape=="circle": lwi=7 else: lwi=8 for ax in axis: ia, view, loc=pms[ax] shp=view.get_shapes() if len(shp)!=len(boxes): ### something changes the box shapes... for i,b in enumerate(boxes): self.update_box_shape(i,b) for ax in axis: ia, view, loc=pms[ax] ## update the box shapes shp=view.get_shapes() for i,b in enumerate(boxes): bs=self.get_boxsize(b[5]) dst=abs(b[ia] - loc) inplane=dst<bs//2 rad=bs//2-dst if ax!='z' and allside: ## display all side view boxes in this mode inplane=True rad=bs//2 if ax=='z' and self.options.mode=="2D": ## boxes are 1 slice thick in 2d mode inplane=dst<1 if inplane and (b[5] in self.sets_visible): shp[i][0]=self.boxshape ## selected box is slightly thicker if self.curbox==i: shp[i][lwi]=3 else: shp[i][lwi]=2 if self.options.mode=="3D": shp[i][6]=rad else: shp[i][0]="hidden" view.shapechange=1 img=self.get_slice(loc, self.nlayers(), ax) if self.wfilt.getValue()!=0.0: img.process_inplace("filter.lowpass.gauss",{"cutoff_freq":1.0/self.wfilt.getValue(),"apix":self.apix}) view.set_data(img) self.update_coords() def get_rotated_boxes(self): if len(self.boxes)==0: return [] if self.globalxf.is_identity(): boxes=self.boxes else: cnt=[self.data["nx"]//2,self.data["ny"]//2,self.data["nz"]//2] pts=np.array([b[:3] for b in self.boxes])-cnt pts=np.array([self.globalxf.transform(p.tolist()) for p in pts])+cnt boxes=[] for i,b in enumerate(self.boxes): p=pts[i] boxes.append([p[0],p[1],p[2],b[3],b[4],b[5]]) return boxes def update_all(self): """redisplay of all widgets""" if self.data==None: return self.update_sliceview() self.update_boximgs() def update_coords(self): self.wcoords.setText("X: {:d}\tY: {:d}\tZ: {:d}".format(int(self.x_loc), int(self.y_loc), int(self.z_loc))) def inside_box(self,n,x=-1,y=-1,z=-1): """Checks to see if a point in image coordinates is inside box number n. If any value is negative, it will not be checked.""" box=self.boxes[n] if box[5] not in self.sets_visible: return False bs=self.get_boxsize(box[5])/2 if self.options.mode=="3D": rr=(x>=0)*((box[0]-x)**2) + (y>=0)*((box[1]-y) **2) + (z>=0)*((box[2]-z)**2) else: rr=(x>=0)*((box[0]-x)**2) + (y>=0)*((box[1]-y) **2) + (z>=0)*(box[2]!=z)*(1e3*bs**2) return rr<=bs**2 def del_box(self, delids): if type(delids)!=list: delids=[delids] kpids=[i for i,b in enumerate(self.boxes) if i not in delids] self.boxes=[self.boxes[i] for i in kpids] self.boxesimgs=[self.boxesimgs[i] for i in kpids] self.xyview.shapes={i:self.xyview.shapes[k] for i,k in enumerate(kpids)} self.xzview.shapes={i:self.xzview.shapes[k] for i,k in enumerate(kpids)} self.zyview.shapes={i:self.zyview.shapes[k] for i,k in enumerate(kpids)} #print self.boxes, self.xyview.get_shapes() self.curbox=-1 self.update_all() def update_box_shape(self,n, box): bs2=self.get_boxsize(box[5])//2 if n==self.curbox: lw=3 else: lw=2 color=self.setcolors[box[5]%len(self.setcolors)].color().getRgbF() if self.options.mode=="3D": self.xyview.add_shape(n,EMShape(["circle",color[0],color[1],color[2],box[0],box[1],bs2,lw])) self.xzview.add_shape(n,EMShape(["circle",color[0],color[1],color[2],box[0],box[2],bs2,lw])) self.zyview.add_shape(n,EMShape(("circle",color[0],color[1],color[2],box[2],box[1],bs2,lw))) else: self.xyview.add_shape(n,EMShape(["rect",color[0],color[1],color[2], box[0]-bs2,box[1]-bs2,box[0]+bs2,box[1]+bs2,2])) self.xzview.add_shape(n,EMShape(["rect",color[0],color[1],color[2], box[0]-bs2,box[2]-1,box[0]+bs2,box[2]+1,2])) self.zyview.add_shape(n,EMShape(["rect",color[0],color[1],color[2], box[2]-1,box[1]-bs2,box[2]+1,box[1]+bs2,2])) def update_box(self,n,quiet=False): """After adjusting a box, call this""" # print "upd ",n,quiet if n<0 or n>=len(self.boxes): return box=self.boxes[n] boxes=self.get_rotated_boxes() self.update_box_shape(n,boxes[n]) if self.initialized: self.update_sliceview() # For speed, we turn off updates while dragging a box around. Quiet is set until the mouse-up if not quiet: # Get the cube from the original data (normalized) proj=self.get_cube(box[0], box[1], box[2], centerslice=True, boxsz=self.get_boxsize(box[5])) proj.process_inplace("normalize") for i in range(len(self.boxesimgs),n+1): self.boxesimgs.append(None) self.boxesimgs[n]=proj mm=[m for im,m in enumerate(self.boxesimgs) if self.boxes[im][5] in self.sets_visible] if self.initialized: self.SaveJson() if self.initialized: self.update_boximgs() #if n!=self.curbox: #self.boxesviewer.set_selected((n,),True) self.curbox=n def update_boximgs(self): self.boxids=[im for im,m in enumerate(self.boxesimgs) if self.boxes[im][5] in self.sets_visible] self.boxesviewer.set_data([self.boxesimgs[i] for i in self.boxids]) self.boxesviewer.update() return def img_selected(self,event,lc): #print("sel",lc[0]) lci=self.boxids[lc[0]] if event.modifiers()&Qt.ShiftModifier: if event.modifiers()&Qt.ControlModifier: self.del_box(list(range(lci, len(self.boxes)))) else: self.del_box(lci) else: #self.update_box(lci) self.curbox=lci box=self.boxes[lci] self.x_loc,self.y_loc,self.z_loc=self.rotate_coord([box[0], box[1], box[2]], inv=False) self.scroll_to(self.x_loc,self.y_loc,self.z_loc) self.update_sliceview() def rotate_coord(self, p, inv=True): if not self.globalxf.is_identity(): cnt=[self.data["nx"]//2,self.data["ny"]//2,self.data["nz"]//2] p=[p[i]-cnt[i] for i in range(3)] xf=Transform(self.globalxf) if inv: xf.invert() p=xf.transform(p)+cnt return p def del_region_xy(self, x=-1, y=-1, z=-1, rad=-1): if rad<0: rad=self.eraser_width() delids=[] boxes=self.get_rotated_boxes() for i,b in enumerate(boxes): if b[5] not in self.sets_visible: continue if (x>=0)*(b[0]-x)**2 + (y>=0)*(b[1]-y)**2 +(z>=0)*(b[2]-z)**2 < rad**2: delids.append(i) self.del_box(delids) def scroll_to(self, x,y,z, axis=""): if axis!="z": self.xyview.scroll_to(x,y,True) if axis!="y": self.xzview.scroll_to(x,self.data["nz"]/2,True) if axis!="x": self.zyview.scroll_to(self.data["nz"]/2,y,True) #### mouse click def xy_down(self,event): x,y=self.xyview.scr_to_img((event.x(),event.y())) self.mouse_down(event, x,y,self.z_loc, "z") def xz_down(self,event): x,z=self.xzview.scr_to_img((event.x(),event.y())) self.mouse_down(event,x,self.y_loc,z, "y") def zy_down(self,event): z,y=self.zyview.scr_to_img((event.x(),event.y())) self.mouse_down(event,self.x_loc,y,z, "x") def mouse_down(self,event, x, y, z, axis): if min(x,y,z)<0: return xr,yr,zr=self.rotate_coord((x,y,z)) #print(x,y,z,xr,yr,zr) if self.optionviewer.erasercheckbox.isChecked(): self.del_region_xy(x,y,z) return for i in range(len(self.boxes)): if self.inside_box(i,xr,yr,zr): if event.modifiers()&Qt.ShiftModifier: ## delete box self.del_box(i) else: ## start dragging self.dragging=i self.curbox=i self.scroll_to(x,y,z,axis) break else: if not event.modifiers()&Qt.ShiftModifier: ## add box self.x_loc, self.y_loc, self.z_loc=x,y,z self.scroll_to(x,y,z,axis) self.curbox=len(self.boxes) self.boxes.append(([xr,yr,zr, 'manual', 0.0, self.currentset])) self.update_box(len(self.boxes)-1) self.dragging=len(self.boxes)-1 #### eraser mode def xy_move(self,event): self.mouse_move(event, self.xyview) def xz_move(self,event): self.mouse_move(event, self.xzview) def zy_move(self,event): self.mouse_move(event, self.zyview) def mouse_move(self,event,view): if self.optionviewer.erasercheckbox.isChecked(): self.xyview.eraser_shape=self.xzview.eraser_shape=self.zyview.eraser_shape=None x,y=view.scr_to_img((event.x(),event.y())) view.eraser_shape=EMShape(["circle",1,1,1,x,y,self.eraser_width(),2]) view.shapechange=1 view.update() else: view.eraser_shape=None #### dragging... def mouse_drag(self,x, y, z): if self.dragging<0: return if min(x,y,z)<0: return self.x_loc, self.y_loc, self.z_loc=x,y,z x,y,z=self.rotate_coord((x,y,z)) self.boxes[self.dragging][:3]= x,y,z self.update_box(self.dragging,True) def xy_drag(self,event): if self.dragging>=0: x,y=self.xyview.scr_to_img((event.x(),event.y())) self.mouse_drag(x,y,self.z_loc) def xz_drag(self,event): if self.dragging>=0: x,z=self.xzview.scr_to_img((event.x(),event.y())) self.mouse_drag(x,self.y_loc,z) def zy_drag(self,event): if self.dragging>=0: z,y=self.zyview.scr_to_img((event.x(),event.y())) self.mouse_drag(self.x_loc,y,z) def mouse_up(self,event): if self.dragging>=0: self.update_box(self.dragging) self.dragging=-1 #### keep the same origin for the 3 views def xy_origin(self,newor): xzo=self.xzview.get_origin() self.xzview.set_origin(newor[0],xzo[1],True) zyo=self.zyview.get_origin() self.zyview.set_origin(zyo[0],newor[1],True) def xz_origin(self,newor): xyo=self.xyview.get_origin() self.xyview.set_origin(newor[0],xyo[1],True) def zy_origin(self,newor): xyo=self.xyview.get_origin() self.xyview.set_origin(xyo[0],newor[1],True) ##### go up/down with shift+wheel def xy_wheel(self, event): z=int(self.z_loc+ np.sign(event.angleDelta().y())) if z>0 and z<self.data["nz"]: self.wdepth.setValue(z) def xz_wheel(self, event): y=int(self.y_loc+np.sign(event.angleDelta().y())) if y>0 and y<self.data["ny"]: self.y_loc=y self.update_sliceview(['y']) def zy_wheel(self, event): x=int(self.x_loc+np.sign(event.angleDelta().y())) if x>0 and x<self.data["nx"]: self.x_loc=x self.update_sliceview(['x']) ######## def set_current_set(self, name): #print "set current", name name=parse_setname(name) self.currentset=name self.wboxsize.setValue(self.get_boxsize()) self.update_all() return def hide_set(self, name): name=parse_setname(name) if name in self.sets_visible: self.sets_visible.pop(name) if self.initialized: self.update_all() self.update_boximgs() return def show_set(self, name): name=parse_setname(name) self.sets_visible[name]=0 #self.currentset=name #self.wboxsize.setValue(self.get_boxsize()) if self.initialized: self.update_all() self.update_boximgs() return def delete_set(self, name): name=parse_setname(name) ## idx to keep delids=[i for i,b in enumerate(self.boxes) if b[5]==int(name)] self.del_box(delids) if name in self.sets_visible: self.sets_visible.pop(name) if name in self.sets: self.sets.pop(name) if name in self.boxsize: self.boxsize.pop(name) self.update_all() return def rename_set(self, oldname, newname): name=parse_setname(oldname) if name in self.sets: self.sets[name]=newname return def new_set(self, name): for i in range(len(self.sets)+1): if i not in self.sets: break self.sets[i]=name self.sets_visible[i]=0 if self.options.mode=="3D": self.boxsize[i]=32 else: self.boxsize[i]=64 return def save_set(self): self.save_boxes(list(self.sets_visible.keys())) return def key_press(self,event): if event.key() == 96: ## "`" to move up a slice since arrow keys are occupied... self.wdepth.setValue(self.z_loc+1) elif event.key() == 49: ## "1" to move down a slice self.wdepth.setValue(self.z_loc-1) else: self.keypress.emit(event) def flatten_tomo(self): print("Flatten tomogram by particles coordinates") vis=list(self.sets_visible.keys()) pts=[b[:3] for b in self.boxes if b[5] in vis] if len(pts)<3: print("Too few visible particles. Cannot flatten tomogram.") return pts=np.array(pts) pca=PCA(3) pca.fit(pts); c=pca.components_ t=Transform() cc=c[2] if cc[2]!=0: cc*=np.sign(cc[2]) t.set_rotation(c[2].tolist()) t.invert() xyz=t.get_params("xyz") xyz["ztilt"]=0 print("xtilt {:.02f}, ytilt {:.02f}".format(xyz["xtilt"], xyz["ytilt"])) t=Transform(xyz) self.globalxf=t self.dataxf=self.data.process("xform",{"transform":t}) self.xyview.shapes={} self.zyview.shapes={} self.xzview.shapes={} boxes=self.get_rotated_boxes() for i,b in enumerate(boxes): self.update_box_shape(i,b) self.update_sliceview() print("Done") def reset_flatten_tomo(self, event): self.globalxf=Transform() self.xyview.shapes={} self.zyview.shapes={} self.xzview.shapes={} boxes=self.get_rotated_boxes() for i,b in enumerate(boxes): self.update_box_shape(i,b) self.update_sliceview() def SaveJson(self): info=js_open_dict(self.jsonfile) sx,sy,sz=(self.data["nx"]//2,self.data["ny"]//2,self.data["nz"]//2) if "apix_unbin" in info: bxs=[] for b0 in self.boxes: b=[ (b0[0]-sx)*self.apix_cur/self.apix_unbin, (b0[1]-sy)*self.apix_cur/self.apix_unbin, (b0[2]-sz)*self.apix_cur/self.apix_unbin, b0[3], b0[4], b0[5] ] bxs.append(b) bxsz={} for k in self.boxsize.keys(): bxsz[k]=np.round(self.boxsize[k]*self.apix_cur/self.apix_unbin) else: bxs=self.boxes bxsz=self.boxsize info["boxes_3d"]=bxs clslst={} for key in list(self.sets.keys()): clslst[int(key)]={ "name":self.sets[key], "boxsize":int(bxsz[key]), } info["class_list"]=clslst info.close() def closeEvent(self,event): print("Exiting") self.SaveJson() E2saveappwin("e2sptboxer","main",self) E2saveappwin("e2sptboxer","boxes",self.boxesviewer.qt_parent) E2saveappwin("e2sptboxer","option",self.optionviewer) #self.boxviewer.close() self.boxesviewer.close() self.optionviewer.close() #self.optionviewer.close() #self.xyview.close() #self.xzview.close() #self.zyview.close() self.module_closed.emit() # this signal is important when e2ctf is being used by a program running its own event loop
class GUIFourierSynth(QtWidgets.QWidget): """This class represents an application for interactive Fourier synthesis""" def __init__(self,app): self.app=app QtWidgets.QWidget.__init__(self,None) self.synthplot=EMPlot2DWidget(self.app) self.synthplot.show() self.fftplot=EMPlot2DWidget(self.app) # not shown initially self.fftplot.show() # self.bispecimg=EMImage2DWidget(self.app) # not shown initially # overall layout self.vbl1=QtWidgets.QVBoxLayout() self.setLayout(self.vbl1) # First row contains general purpose controls self.hbl1=QtWidgets.QHBoxLayout() self.vbl1.addLayout(self.hbl1) self.vcell=ValBox(self,(0,128.0),"Cell:",64) self.hbl1.addWidget(self.vcell) self.vncells=ValBox(self,(0,128.0),"n Cells:",1) self.hbl1.addWidget(self.vncells) self.voversamp=ValBox(self,(0,128.0),"Oversample:",1) self.hbl1.addWidget(self.voversamp) self.targfn=None self.vnsin=ValBox(self,(1,64),"# Sin:",32) self.vnsin.intonly=1 self.hbl1.addWidget(self.vnsin) self.cbshowall=QtWidgets.QCheckBox("Show All") self.hbl1.addWidget(self.cbshowall) self.cbshifted=QtWidgets.QCheckBox("Shifted") self.hbl1.addWidget(self.cbshifted) self.vshftstep=ValBox(self,(1,64),"Shft:",1) self.hbl1.addWidget(self.vshftstep) self.bphaseleft=QtWidgets.QPushButton("\u2190") # phase left self.hbl1.addWidget(self.bphaseleft) self.bphasecen=QtWidgets.QPushButton("O") self.hbl1.addWidget(self.bphasecen) self.bphaseright=QtWidgets.QPushButton("\u2192") # phase right self.hbl1.addWidget(self.bphaseright) self.cbtargfn=QtWidgets.QComboBox(self) self.cbtargfn.addItem("None") # 0 self.cbtargfn.addItem("triangle") # 1 self.cbtargfn.addItem("square") # 2 self.cbtargfn.addItem("square imp") # 3 self.cbtargfn.addItem("delta") # 4 self.cbtargfn.addItem("noise") self.cbtargfn.addItem("saw") self.cbtargfn.addItem("sin") self.cbtargfn.addItem("modsin") self.cbtargfn.addItem("modsin2") self.cbtargfn.addItem("modsin3") #10 self.cbtargfn.addItem("sin low") self.cbtargfn.addItem("doubledelta") self.cbtargfn.addItem("sin bad f") self.cbtargfn.addItem("sin bad f2") self.cbtargfn.addItem("0 phase") #15 self.cbtargfn.addItem("rand phase") self.hbl1.addWidget(self.cbtargfn) self.bsound=QtWidgets.QPushButton("Play") self.hbl1.addWidget(self.bsound) self.bsoundr=QtWidgets.QPushButton("Rec") self.hbl1.addWidget(self.bsoundr) self.vrootf=ValBox(self,(1,64),"Root F:",264) # 264 is middle C #self.vrootf.intonly=1 self.hbl1.addWidget(self.vrootf) # Widget containing valsliders self.wapsliders=QtWidgets.QWidget(self) # self.wapsliders.setMinimumSize(800,640) self.gblap=QtWidgets.QGridLayout() self.gblap.setSizeConstraint(QtWidgets.QLayout.SetMinAndMaxSize) self.gblap.setColumnMinimumWidth(0,250) self.gblap.setColumnMinimumWidth(1,250) self.wapsliders.setLayout(self.gblap) # ScrollArea providing view on slider container widget self.wapsarea=QtWidgets.QScrollArea(self) self.wapsarea.setWidgetResizable(True) self.wapsarea.setWidget(self.wapsliders) self.vbl1.addWidget(self.wapsarea) self.vcell.valueChanged.connect(self.recompute) self.vncells.valueChanged.connect(self.recompute) self.voversamp.valueChanged.connect(self.recompute) self.vnsin.valueChanged.connect(self.nsinchange) self.cbshowall.stateChanged[int].connect(self.recompute) self.cbshifted.stateChanged[int].connect(self.recompute) self.cbtargfn.activated[int].connect(self.newtargfn) self.bphaseleft.clicked.connect(self.phaseleft) self.bphasecen.clicked.connect(self.phasecen) self.bphaseright.clicked.connect(self.phaseright) self.bsound.clicked.connect(self.playsound) self.bsoundr.clicked.connect(self.recsound) self.wamp=[] self.wpha=[] self.curves=[] self.xvals=[] for i in range(65): self.wamp.append(ValSlider(self,(0.0,1.0),"%2d:"%i,0.0)) self.gblap.addWidget(self.wamp[-1],i,0) self.wamp[-1].valueChanged.connect(self.recompute) self.wpha.append(ValSlider(self,(-180.0,180.0),"%2d:"%i,0.0)) self.gblap.addWidget(self.wpha[-1],i,1) self.wpha[-1].valueChanged.connect(self.recompute) self.curves.append(EMData(64,1)) # if self.cbshowall.isChecked() : # self.synthplot self.total=EMData(64,1) self.nsinchange() E2loadappwin("e2fftsynth","main",self) E2loadappwin("e2fftsynth","synth",self.synthplot.qt_parent) E2loadappwin("e2fftsynth","fft",self.fftplot.qt_parent) def closeEvent(self,event): # QtWidgets.QWidget.closeEvent(self,event) E2saveappwin("e2fftsynth","main",self) E2saveappwin("e2fftsynth","synth",self.synthplot.qt_parent) E2saveappwin("e2fftsynth","fft",self.fftplot.qt_parent) QtWidgets.qApp.exit(0) def phaseleft(self,v): # fixed translation in minus direction ss=self.vshftstep.getValue() sft=ss*360.0/len(self.xvals) for i in range(1,len(self.wpha)): self.wpha[i].setValue(fixang(self.wpha[i].getValue()-sft*i),1) self.recompute() def phasecen(self,v): # translation so phase of lowest freqeuncy is zero sft=self.wpha[1].getValue() for i in range(1,len(self.wpha)): self.wpha[i].setValue(fixang(self.wpha[i].getValue()-sft*i),1) self.recompute() def phaseright(self,v): # fixed translation in plus direction ss=self.vshftstep.getValue() sft=ss*360.0/len(self.xvals) for i in range(1,len(self.wpha)): self.wpha[i].setValue(fixang(self.wpha[i].getValue()+sft*i),1) self.recompute() def playsound(self,v): global soundout f0=int(self.vrootf.getValue()) nsam2=8*44100//(2*f0) # 8x oversampling so high fundamental frequencies come out closer to correct print(f"play: {f0} ({nsam2})") svals=np.zeros_like(self.svals,shape=nsam2) if len(self.svals)>=nsam2//8: svals[::8]=self.svals[:(nsam2-1)//8+1] # interleave zeros for 8 repetitions of the wave. Gives better frequency accuracy at high frequency else: svals[:len(self.svals)*8:8]=self.svals # svals=self.svals.copy() # svals.resize(nsam2) # if len(self.svals)<nsam2: svals[len(self.svals):]=0.0 # zero fill if we extended the array total=fft.irfft(svals)*(len(svals)-1)*10000.0 self.assound=np.tile(total.astype("int16"),44100//nsam2) #print nsam2,len(self.svals),len(svals),len(total) initsound() soundout.start() soundout.write(self.assound) soundout.stop() def recsound(self,v): global soundin # global kaiser initsound() soundin.start() snd=soundin.read(65536)[0][:,0] soundin.stop() np.savetxt("real.txt",snd) # snd*=kaiser # window didn't really help ft=fft.rfft(snd) a=np.absolute(ft) np.savetxt("fft.txt",a) p=np.angle(ft)*180./pi np.savetxt("pha.txt",p) # If there is only a single strong peak (or maybe 2) the CCF approach won't work very well, but the peak itself should be good # h=np.argwhere(a>np.max(a)/20.0) # if len(h)<3: h=np.min(h) # else: fta=fft.rfft(a) fta=fta*np.conj(fta) fta=fft.irfft(fta) np.savetxt("fta.txt",fta) # peak filter, but wasn't useful # ftap=np.roll(fta,1,0) # ftam=np.roll(fta,-1,0) # fta=np.where(fta>ftap,fta,np.zeros(len(fta))) # fta=np.where(fta>ftam,fta,np.zeros(len(fta))) # np.savetxt("ftaf.txt",fta) # ftac=fta.copy() # # need a clever way to do this in numpy # for i in range(1,32767): # if fta[i-1]>fta[i] or fta[i+1]>fta[i] : ftac[i]=0 h=(np.argmax(fta[200:20000])+200)//2 f1=22050.0/32768.0 # lowest frequency pixel, scaling factor for frequency axis # sum the first 8 harmonics for each fundamental # hsum=a[:4096]*1.25 # *1.25 helps make sure a solo peak shows up as the first harmonic not a higher one # for i in range(2,9): hsum+=a[:4096*i:i] # hsum=np.fmin(hsum,a[:4096]*25.0) # np.savetxt("fsum.txt",hsum) # hsum=hsum[100:] # h=np.min(np.argwhere(hsum>np.std(hsum)*2.0))+100 # h=np.argmax(hsum[100:])+100 maxf=f1*h # should correspond to the strongest single frequency based on harmonic sum print(f"peak {h} -> {maxf}") self.vrootf.setValue(floor(maxf)) # pull out the values for up to the first 33 harmonics sca=0.5/np.max(a[h::h]) for i in range(1,33): j=h*i # harmonic peak index (rounded) if j>32767 : self.wamp[i].setValue(0,True) continue amp=float(abs(ft[j]))*sca pha=float(cmath.phase(ft[j]))*180.0/pi self.wamp[i].setValue(amp,True) self.wpha[i].setValue(pha,True) self.recompute() def newtargfn(self,index): "This function has a number of hardcoded 'target' functions" if index==0 : self.targfn=None nsin=int(self.vnsin.getValue()) for i in range(nsin+1): self.wamp[i].setValue(0,1) self.wpha[i].setValue(0,1) elif index==15: # zero phase nsin=int(self.vnsin.getValue()) for i in range(nsin+1): self.wpha[i].setValue(0,1) elif index==16: # random phase nsin=int(self.vnsin.getValue()) for i in range(nsin+1): self.wpha[i].setValue(random.uniform(-180.,180.),1) else : nx=int(self.vcell.getValue()) x=np.arange(0,1.0,1.0/nx) y=np.zeros(nx) if index==1 : # triangle y[:nx//2]=x[:nx//2]*2.0 y[nx//2:]=(1.0-x[nx//2:])*2.0 elif index==2 : # square y[:nx//4]=0 y[nx//4:nx*3//4]=1.0 y[nx*3//4:]=0 elif index==3 : # square impulse y[:nx//2]=0 y[nx//2:nx*3//5]=1.0 y[nx*3//5:]=0 elif index==4 : # delta y[:]=0 y[nx//2]=10.0 elif index==5 : # noise y=np.random.random_sample((nx,))-0.5 elif index==6 : # saw y[:nx//4]=0 y[nx//4:nx//2]=x[:nx//4] y[nx//2:nx*3//4]=-x[:nx//4] y[nx*3//4:]=0 elif index==7 : # sin y=np.sin(x*16.0*pi) elif index==8 : # modulated sine y=np.sin(x*4.0*pi)*np.sin(x*32.0*pi) elif index==9 : # modulated sine 2 y=np.sin(x*2.0*pi)*np.sin(x*32.0*pi) elif index==10 : # modulated sine 3 y=np.sin(x*8.0*pi)*np.sin(x*32.0*pi) elif index==11 : # sin low y=np.sin(x*4.0*pi) elif index==12 : # double delta y[:]=0 y[16]=5.0 y[48]=5.0 elif index==13 : # sin bad f y=np.sin(x*2.3*pi) elif index==14 : # sin bad f2 y=np.sin(x*17.15*pi) self.targfn=y self.target2sliders() self.recompute() def target2sliders(self): f=fft.rfft(self.targfn)*2.0/len(self.targfn) nsin=int(self.vnsin.getValue()) for i in range(min(len(f),nsin+1)): # print fft[i] amp=abs(f[i]) if fabs(amp)<1.0e-6 : amp=0.0 if amp==0 : pha=0 else: pha=cmath.phase(f[i]) self.wamp[i].setValue(amp,quiet=1) self.wpha[i].setValue(fixang(pha*180.0/pi),quiet=1) def nsinchange(self,value=None): if value==None : value=int(self.vnsin.getValue()) for i in range(65): if i>value: self.wamp[i].hide() self.wpha[i].hide() else : self.wamp[i].show() self.wpha[i].show() if self.targfn!=None : self.target2sliders() self.recompute() def recompute(self,value=None): nsin=int(self.vnsin.getValue()) cell=int(self.vcell.getValue()) ncells=int(self.vncells.getValue()) oversamp=max(1,int(self.voversamp.getValue())) samples=int(cell*oversamp) # arrays since we're using them several ways self.svals=np.array([cmath.rect(self.wamp[i].getValue(),self.wpha[i].getValue()*pi/180.0) for i in range(nsin+1)]) #self.wamps=np.array([v.getValue() for v in self.wamp[:nsin+1]]) #self.wphas=np.array([v.getValue() for v in self.wpha[:nsin+1]]) self.xvals=np.array([xn/float(oversamp) for xn in range(samples)]) if samples//2>len(self.svals): svals=np.concatenate((self.svals,np.zeros(1+samples//2-len(self.svals)))) else: svals=self.svals self.total=fft.irfft(svals)*(len(svals)-1) if ncells>1: self.total=np.tile(self.total,ncells) self.synthplot.set_data((np.arange(0,len(self.total)/oversamp,1.0/oversamp),self.total),"Sum",replace=True,quiet=True,linewidth=2) if not self.targfn is None: self.synthplot.set_data((np.arange(len(self.targfn)),self.targfn),"Target",quiet=True,linewidth=1,linetype=2,symtype=0) # if self.cbshowall.isChecked() : # csum=self.total["minimum"]*1.1 # for i in range(nsin): # if self.wamps[i]==0: continue # csum-=self.wamps[i]*1.1 # if self.cbshifted.isChecked() : self.curves[i].add(csum) # self.synthplot.set_data((self.xvals,self.curves[i].get_data_as_vector()),"%d"%i,quiet=True,linewidth=1,color=2) # csum-=self.wamps[i]*1.1 self.synthplot.updateGL() self.fftplot.set_data((np.arange(len(self.svals)),np.abs(self.svals)),"Amp",color=0,linetype=0,linewidth=2) self.fftplot.set_data((np.arange(len(self.svals)),np.angle(self.svals)),"Pha",color=1,linetype=0,linewidth=2) self.fftplot.updateGL()