Add files via upload

classes/Dataset.py

@@ -0,0 +1,141 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Fri Aug  3 13:29:11 2018
+
+@author: Yi Zhang
+@reference: Deep Unfolded Robust PCA With Application to Clutter Suppression in Ultrasound. https://github.com/KrakenLeaf/CORONA
+
+"""
+
+import numpy as np
+import os
+import torch
+from scipy.io import loadmat
+import random
+
+class Dataset:
+    def __init__(self,folder=None,shuffle=None,prefix=None,shownum=10):
+        self._folder=None
+        if not folder is None:
+            self._folder=folder
+            self._flist=os.listdir(folder)
+        if not prefix is None:    
+            tmp=[]
+            for s in self._flist:
+                if s.startswith(prefix):
+                    tmp.append(s)
+            self._flist=tmp
+        self._flist.sort()
+        if not shuffle is None:
+            random.shuffle(self._flist,shuffle)
+            print(self._flist[0:shownum])
+
+    def getnpz(self,fileind,arrind,folder=None):
+        if folder is None:
+            folder=self._folder
+            flist=self._flist
+        else:
+            flist=os.listdir(folder)
+        data=np.load(folder+flist[fileind],mmap_mode='r+')
+        dlist=[]
+        l=0
+        for x in data:
+            if l in arrind:
+                dlist.append(data[x])
+            l=l+1
+        return dlist
+
+    def getmat(self,fileind,arrind,folder=None):
+        if folder is None:
+            folder=self._folder
+            flist=self._flist
+        else:
+            flist=os.listdir(folder)
+        data=loadmat(folder+flist[fileind])
+        dlist=[]
+        for x in arrind:
+            dlist.append(data[x])
+        return dlist
+
+
+class Converter:
+    def __init__(self):
+        """
+        Preprocessing:
+        1.concat: inv=False, concatenate real and imaginary parts in axis=-1
+            inv=True, depart two parts in -1 axis into real and imaginary parts
+        2.ch2: inv=False, concatenate real and imaginary parts in a new axis=0
+            inv=True, depart two parts in 1 axis into real and imaginary parts
+        3.stack: inv=False, stack real and imaginary parts in axis=-1
+            inv=True, depart two parts in -1 axis into real and imaginary parts
+        3.None: pass
+        """        
+        self.pre={'concat':self.concat,'ch2':self.ch2,'stack':self.stack}
+
+    def np2torch(self,xnp,formlist):
+        dlist=[]
+        for i,x in enumerate(xnp):
+            if ('pre' in formlist[i]) and (formlist[i]['pre'] in self.pre):
+                x=self.pre[formlist[i]['pre']](x,inv=False)
+            x=x.reshape(formlist[i]['shape'])
+            x=torch.tensor(x,dtype=torch.float32)
+            dlist.append(x)
+        return dlist
+
+    def torch2np(self,xtorch,formlist):
+        dlist=[]
+        for i,x in enumerate(xtorch):
+            if torch.cuda.is_available():
+                x=x.cpu().detach().numpy()
+            else:
+                x=x.detach().numpy()
+            if ('pre' in formlist[i]) and (formlist[i]['pre'] in self.pre):
+                x=self.pre[formlist[i]['pre']](x,inv=True)
+            x=x.reshape(formlist[i]['shape'])            
+            dlist.append(x)
+        return dlist
+
+    def stack(self,x,inv):
+        if not inv:
+            #not inv, np2torch
+            x=np.swapaxes(x,0,-1)
+            xr=x.real
+            xi=x.imag
+            size=list(x.shape)
+            size[0]*=2
+            size=tuple(size)
+            z=np.zeros(size)
+            z[0::2]=xr
+            z[1::2]=xi
+            z=np.swapaxes(z,0,-1)
+        else:
+            #inv, torch2np
+            #size=[B,C,H,W,(T)],numpy
+            x=np.swapaxes(x,0,-1)
+            xr=x[0::2]
+            xi=x[1::2]
+            z=xr+1j*xi
+            z=np.swapaxes(z,0,-1)
+
+        return z
+
+    def concat(self,x,inv):
+        if not inv:
+            #not inv, np2torch
+            z=np.concatenate((x.real,x.imag),axis=-1)
+        else:
+            #inv, torch2np
+            #size=[B,C,H,W,(T)],numpy
+            x=np.swapaxes(x,0,-1)
+            n=x.shape[0]
+            nh=int(n/2)
+            xr=x[0:nh]
+            xi=x[nh:n]
+            z=xr+1j*xi
+            z=np.swapaxes(z,0,-1)
+
+        return z
+
+    def ch2(self,x,inv):
+        pass
+

classes/Player.py

@@ -0,0 +1,114 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Thu Aug  2 15:43:10 2018
+
+@author: Yi Zhang
+@reference: Deep Unfolded Robust PCA With Application to Clutter Suppression in Ultrasound. https://github.com/KrakenLeaf/CORONA
+
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import torch
+
+minDBdf=-50
+class Player:
+    def __init__(self,Tpause=None):
+        if Tpause==None:
+            self.Tpause=0.1
+        else:
+            self.Tpause=Tpause
+        self.fig=None
+        self.ax={1:None,2:None,3:None,4:None,5:None,6:None}
+        self.axnum=0
+
+    def plotmat(self,mvlist,note=None,tit=None,supt=None,
+                cmap='gray',ion=True,minDB=None):
+        """
+        input:matrix dimension
+        """
+        if minDB is None:
+            minDB=minDBdf
+
+        if ion:
+            plt.ion()
+        subp={1:[1,1],2:[1,2],3:[1,3],4:[2,2],5:[2,3],6:[2,3],9:[3,3],12:[4,3]}
+        p1,p2=subp[len(mvlist)]
+        if p1*p2!=self.axnum or self.fig is None\
+           or not(plt.fignum_exists(self.fig.number)):
+            self.fig,(self.ax)=plt.subplots(p1,p2) 
+            self.axnum=p1*p2
+            if self.axnum==1:
+                self.ax=np.array([self.ax])   
+            self.ax=self.ax.reshape([-1])
+
+        for i in range(len(mvlist)):            
+            US=mvlist[i]
+            if US is None:
+                continue
+            if US.dtype is torch.float32:
+                US=US.detach().numpy()
+            US=np.abs(US).reshape([-1,mvlist[i].shape[-1]])
+            if np.sum(np.abs(US))!=0:
+                US=US/np.max(US)
+            if note=='db':
+                US[US<10**(minDB/20)]=10**(minDB/20)
+                US=20*np.log10(US)
+            vmin,vmax=[minDB,0] if note=='db' else [0,1]
+            self.ax[i].clear()
+            self.ax[i].imshow(US,cmap=cmap,aspect='auto',vmin=vmin,vmax=vmax)
+            if not(tit is None):
+                self.ax[i].set_title(tit[i])
+        if not(supt is None):
+            self.fig.suptitle(supt)
+        if ion:
+            plt.pause(self.Tpause)
+        return self.fig
+
+    def play(self,mvlist,note=None,tit=None,supt=None,cmap='gray',minDB=None):
+        """
+        input:movie dimension
+        """
+        if minDB is None:
+            minDB=minDBdf
+
+        subp={1:[1,1],2:[1,2],3:[1,3],4:[2,2],5:[2,3],6:[2,3],9:[3,3]}
+        p1,p2=subp[len(mvlist)]
+        T=mvlist[0].shape[-1]
+        fig,ax=plt.subplots(p1,p2) 
+        if p1*p2==1:
+            ax=np.array([ax]) 
+        ax=ax.reshape([-1])
+
+        plt.ion()
+
+        for i in range(len(mvlist)):
+            US=mvlist[i]
+            if US is None:
+                continue
+            if US.dtype is torch.float32:
+                US=US.detach().numpy().squeeze()
+
+            US=np.abs(US)
+            if np.sum(np.abs(US))!=0:
+                US=US/np.max(US)
+            if note=='db':
+                US[US<10**(minDB/20)]=10**(minDB/20)
+                US=20*np.log10(US)
+            mvlist[i]=US
+
+        for t in range(T):
+            for i in range(len(mvlist)):  
+                if mvlist[i] is None:
+                    continue
+                vmin,vmax=[minDB,0] if note=='db' else [0,1]
+                ax[i].clear()
+                ax[i].imshow(mvlist[i][:,:,t],cmap=cmap,aspect='auto',
+                             vmin=vmin,vmax=vmax)
+                if not(tit is None):
+                    ax[i].set_title(tit[i])
+            if supt==None:
+                supt=''
+            fig.suptitle('%dth Frame,'%(t+1)+supt)
+            plt.pause(self.Tpause)
+

classes/Preprocessor.py

@@ -0,0 +1,96 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Wed Aug  8 13:17:20 2018
+
+@author: Yi Zhang
+@reference: Deep Unfolded Robust PCA With Application to Clutter Suppression in Ultrasound. https://github.com/KrakenLeaf/CORONA
+
+"""
+
+import numpy as np
+
+class Preprocessor:
+    def cSVT(self,D,rank,getS1=False):
+        """
+        S: output of cSVT
+        S1: complementary of S
+        """
+        u,s,vh=np.linalg.svd(D)
+        s1=s[0:rank]
+        s[0:rank]=0
+        m,n=D.shape
+        S=np.zeros(D.shape,dtype=np.complex128)
+
+        for i in np.arange(rank,min(m,n)):
+            S[:,i]=u[:,i]*s[i]
+        S=S@vh
+
+        if getS1:
+            S1=np.zeros(D.shape,dtype=np.complex128)
+            for i in range(rank):
+                S1[:,i]=u[:,i]*s1[i]
+            S1=S1@vh
+            return S,S1
+
+        return S
+
+    def cSVT1(self,D):
+        shape=D.shape
+        D=D.reshape([shape[0]*shape[1],shape[2]])
+        u,s,vh=np.linalg.svd(D)
+        s[0]=0
+        m,n=D.shape
+        S=np.zeros(D.shape,dtype=np.complex128)
+
+        for i in np.arange(1,min(m,n)):
+            S[:,i]=u[:,i]*s[i]
+        S=S@vh
+
+        return S.reshape(shape)
+
+    def block_process(self,data,proc,params):
+        """
+        stride should be less than half of block size
+        drop should be less than half of block size
+        """
+        shape=np.array(data.shape)
+        Out=np.zeros(shape,dtype=data.dtype)
+        bshape=np.array(params['bshape'])
+        dim=len(bshape)
+        stride=np.array(params['stride'])
+        drop=np.array(params['drop'])
+        bbeg=np.zeros([dim]).astype(np.int)
+        bend=np.zeros([dim]).astype(np.int)+np.array(bshape)
+
+        while(True):
+            #drop part of the block
+            bend=np.minimum(shape,bend)
+            dropbeg=(bbeg!=0).astype(np.int)
+            dropend=(bend!=shape).astype(np.int)
+            beg=bbeg+dropbeg*drop
+            end=bend-dropend*drop
+            br=beg-bbeg
+            er=end-bbeg
+
+            if dim==2:
+                tmp=proc(data[bbeg[0]:bend[0],bbeg[1]:bend[1]])
+                Out[beg[0]:end[0],beg[1]:end[1]]=tmp[br[0]:er[0],br[1]:er[1]]
+                #print(tmp[br[0]:er[0],br[1]:er[1]])
+                #print(Out)
+            elif dim==3:
+                tmp=proc(data[bbeg[0]:bend[0],bbeg[1]:bend[1],bbeg[2]:bend[2]])
+                Out[beg[0]:end[0],beg[1]:end[1],beg[2]:end[2]]=tmp[br[0]:er[0],br[1]:er[1],br[2]:er[2]]
+
+            #change the block
+            for ii in range(dim):
+                bbeg[ii]+=stride[ii]
+                if (bbeg<shape).all() and  not(bend[ii]==shape[ii]):
+                    break
+                bbeg[0:ii+1]=0
+
+            bend=bbeg+bshape
+            if not bbeg.any():
+                break
+
+        return Out
+

classes/__init__.py

@@ -0,0 +1 @@
+#