updates for PR reviews

diffraction/WISH/bragg-detect/bragg_utils.py

@@ -36,14 +36,14 @@ def createPeaksWorkspaceFromIndices(ws, peak_wsname, indices, data):
     """
     ispec = findSpectrumIndex(indices, *data.shape[0:2])
     peaks = CreatePeaksWorkspace(InstrumentWorkspace=ws, NumberOfPeaks=0, 
-                OutputWorkspace=peak_wsname)
+                OutputWorkspace=peak_wsname, EnableLogging=False)
     for ipk in range(len(ispec)):
         wavelength = ws.readX(ispec[ipk])[indices[ipk,2]]
         # four detectors per spectrum so use one of the central ones
         detIDs = ws.getSpectrum(ispec[ipk]).getDetectorIDs()
         idet = (len(detIDs)-1)//2  # pick central pixel
         AddPeak(PeaksWorkspace=peaks, RunWorkspace=ws, TOF=wavelength, DetectorID=detIDs[idet],
-            BinCount=data[indices[ipk,0], indices[ipk,1], indices[ipk,2]])
+            BinCount=data[indices[ipk,0], indices[ipk,1], indices[ipk,2]], EnableLogging=False)
     return peaks
 
 

diffraction/WISH/bragg-detect/cnn/WISHMLBraggPeaksDetector.py → diffraction/WISH/bragg-detect/cnn/BraggDetectCNN.py

@@ -9,7 +9,19 @@
 from Diffraction.single_crystal.base_sx import BaseSX
 import time
 
-class WISHMLBraggPeaksDetector:
+class BraggDetectCNN:
+    """
+    Detects Bragg's peaks from a workspace using a pre trained deep learning model created using Faster R-CNN model with a ResNet-50-FPN backbone.
+    Example usage is as shown below.
+
+    #1) Set the path of the .pt file that contains the weights of the pre trained FasterRCNN model
+    cnn_weights_path = r""
+
+    # 2) Create a peaks workspace containing bragg peaks detected with a confidence greater than conf_threshold
+    cnn_bragg_peaks_detector = BraggDetectCNN(model_weights_path=cnn_weights_path, batch_size=64, workers=0, iou_threshold=0.001)
+    cnn_bragg_peaks_detector.find_bragg_peaks(workspace="WISH00042730", conf_threshold=0.0, q_tol=0.05)
+    """
+
     def __init__(self, model_weights_path, batch_size=64, workers=0, iou_threshold=0.001):
         """
         :param model_weights_path: Path to the .pt file containing the weights of the pre trained CNN model
@@ -25,20 +37,18 @@ def __init__(self, model_weights_path, batch_size=64, workers=0, iou_threshold=0
         self.iou_threshold = iou_threshold
 
 
-    def find_bragg_peaks(self, run_name, output_ws_name="CNN_Peaks", conf_threshold=0.0, q_tol=0.05):
+    def find_bragg_peaks(self, workspace, output_ws_name="CNN_Peaks", conf_threshold=0.0, q_tol=0.05):
         """
         Find bragg peaks using the pre trained FasterRCNN model and create a peaks workspace
-        :param run_name: Name of the WISH run ex: WISH0042730
+        :param workspace: Workspace name or the object of Workspace from WISH, ex: "WISH0042730"
         :param output_ws_name: Name of the peaks workspace
         :param conf_threshold: Confidence threshold to filter peaks inferred from RCNN
         :param q_tol: qlab tolerance to remove duplicate peaks
         """
         start_time = time.time()
-        data_set, predicted_indices = self._do_cnn_inferencing(run_name)
-        print(f"Predicted indices {predicted_indices}")
+        data_set, predicted_indices = self._do_cnn_inferencing(workspace)
         filtered_indices = predicted_indices[predicted_indices[:, -1] > conf_threshold]
         filtered_indices_rounded = np.round(filtered_indices[:, :-1]).astype(int)
-        print("Creating peaks worspace")
         peaksws = createPeaksWorkspaceFromIndices(data_set.get_workspace(), output_ws_name, filtered_indices_rounded, data_set.get_ws_as_3d_array())
         for ipk, pk in enumerate(peaksws):
             pk.setIntensity(filtered_indices[ipk, -1])
@@ -48,9 +58,8 @@ def find_bragg_peaks(self, run_name, output_ws_name="CNN_Peaks", conf_threshold=
         print(f"Bragg peaks finding from FasterRCNN model is completed in {time.time()-start_time} seconds!")
 
 
-    def _do_cnn_inferencing(self, ws_name):
-        print(f"Do CNN inferencing for {ws_name}")
-        data_set = WISHWorkspaceDataSet(ws_name)
+    def _do_cnn_inferencing(self, workspace):
+        data_set = WISHWorkspaceDataSet(workspace)
         data_loader = tc.utils.data.DataLoader(data_set, batch_size=self.batch_size, shuffle=False, num_workers=self.workers)
         self.model.eval()
         predicted_indices_with_score = []
@@ -98,11 +107,4 @@ def _get_fasterrcnn_resnet50_fpn(self, num_classes=2):
         # replace the head
         model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes) 
         return model
-
-if __name__ == "__main__" or __name__ == "mantidqt.widgets.codeeditor.execution":
-    #1) Set the path of the .pt file that contains the weights of the pre trained FasterRCNN model
-    cnn_weights_path = r""
-
-    # 2) Create a peaks workspace containing bragg peaks detected with a confidence greater than conf_threshold
-    cnn_bragg_peaks_detector = WISHMLBraggPeaksDetector(model_weights_path=cnn_weights_path, batch_size=64, workers=0, iou_threshold=0.001)
-    cnn_bragg_peaks_detector.find_bragg_peaks(ws_name="WISH00042730", conf_threshold=0.0, q_tol=0.05)
+

diffraction/WISH/bragg-detect/cnn/README.md

@@ -1,21 +1,23 @@
 Bragg Peaks detection using a Faster RCNN model
 ================
 
-Inorder to use the pretrained Faster RCNN model inside mantid below steps are required.
+Inorder to use the pretrained Faster RCNN model inside mantid, below steps are required.
 
 * Install mantid from conda `mamba create -n mantid_cnn -c mantid mantidworkbench`
 * Activate the conda environment with `mamba activate mantid_cnn` 
 * Launch workbench from `workbench` command
-* Download the script repository's `scriptrepository\diffraction\WISH` directory as instructed here https://docs.mantidproject.org/nightly/workbench/scriptrepository.html 
+* Download the script repository's `scriptrepository\diffraction\WISH` directory as instructed here https://docs.mantidproject.org/nightly/workbench/scriptrepository.html
+* Check whether `<local path>\diffraction\WISH` path is available at `Python Script Directories` tab from `File->Manage User Directories`.
 * Close the workbench
-* From command line, change the directory to the place where the scripts were downloaded ex: LocalScriptRepo\diffraction\WISH 
+* From command line, change the directory to the place where the scripts were downloaded ex: `<local path>\diffraction\WISH`
 * Within the same conda enviroment, install pytorch dependancies by running `pip install -r requirements.txt`
 * Install NVIDIA CUDA Deep Neural Network library (cuDNN) by running `conda install -c anaconda cudnn`
 * Re-launch workbench from `workbench` command
 * Below is an example code snippet to test the code. It will create a peaks workspace with the inferred peaks from the cnn and will do a peak filtering using the q_tol provided using `BaseSX.remove_duplicate_peaks_by_qlab`.
 ```python
-from cnn.WISHMLBraggPeaksDetector import WISHMLBraggPeaksDetector
+from cnn.BraggDetectCNN import BraggDetectCNN
 model_weights = r'path/to/pretrained/fasterrcnn_resnet50_model_weights.pt'
-cnn_peaks_detector = WISHMLBraggPeaksDetector(model_weights_path=model_weights, batch_size=64)
-cnn_peaks_detector.find_bragg_peaks(run_name='WISH00042730', q_tol=0.05)
+cnn_peaks_detector = BraggDetectCNN(model_weights_path=model_weights, batch_size=64)
+cnn_peaks_detector.find_bragg_peaks(workspace='WISH00042730', output_ws_name="CNN_Peaks", conf_threshold=0.0, q_tol=0.05)
 ```
+* If the above import is not working, check whether the `<local path>\diffraction\WISH` path is listed under `Python Script Directories` tab from `File->Manage User Directories`.

diffraction/WISH/bragg-detect/cnn/WISHDataSets.py

@@ -3,20 +3,27 @@
 import albumentations as A
 from albumentations.pytorch.transforms import ToTensorV2
 from bragg_utils import make_3d_array
-from mantid.simpleapi import Load, Rebunch
+from mantid.simpleapi import Load, Rebunch, Workspace
 
 
 class WISHWorkspaceDataSet(tc.utils.data.Dataset):
-    def __init__(self, ws_name):
-        ws = Load(Filename = ws_name, OutputWorkspace=ws_name)
-        self.rebunched_ws = Rebunch(InputWorkspace=ws, NBunch=3, OutputWorkspace=ws_name+"_rebunched")
+    def __init__(self, workspace):
+        if isinstance(workspace, Workspace):
+            if workspace.getAxis(0).getUnit().unitID() != "TOF":
+                raise RuntimeError("Unit of the X-axis is expected to be TOF")
+            ws = workspace
+        elif isinstance(workspace, str):
+            ws = Load(Filename=workspace, OutputWorkspace=workspace, EnableLogging=False)
+        else:
+            raise RuntimeError("Invalid workspace type - must be Workspace object or a name of a workspace to Load")
+        self.rebunched_ws = Rebunch(InputWorkspace=ws, NBunch=3, OutputWorkspace="_cnn_rebunched", StoreInADS=False, EnableLogging=False)
         self.ws_3d = make_3d_array(self.rebunched_ws)
-        print(f"Data set for {ws_name} is created with shape{self.ws_3d.shape}")
+        print(f"Data set for {workspace} is created with shape{self.ws_3d.shape}")
         self.trans = A.Compose([A.pytorch.transforms.ToTensorV2(p=1.0)])
 
     def get_workspace(self):
         return self.rebunched_ws
-
+        
     def get_ws_as_3d_array(self):
         return self.ws_3d
 

diffraction/WISH/bragg-detect/process_input_and_output_WISH_ML_peak_finding.py

@@ -4,7 +4,7 @@
 import h5py
 from os import path
 from bragg_detect import detect_bragg_peaks
-
+from bragg_utils import make_3d_array, createPeaksWorkspaceFromIndices
 
 # functions to process data to use as ML input
 
@@ -22,26 +22,6 @@ def load_raw_run(runno):
     ws = ConvertToPointData(InputWorkspace=str(runno), OutputWorkspace=str(runno))
     return ws
 
-def make_3d_array(ws, ntubes=1520, npix_per_tube=128, ntubes_per_bank=152, nmonitors=5):
-    """
-    Extract Ydata from WISH MatrixWorkspace and shape into 3d array
-    :param ws: MatrixWorkspace of WISH data with xunit wavelength
-    :param ntubes: number of tubes in instrument (WISH)
-    :param npix_per_tube: number of detector pixels in each tube
-    :param ntubes_per_bank: number of tubes per bank
-    :param nmonitors: number of monitor spectra (assumed first spectra in ws)
-    :return: 3d numpy array ntubes x npix per tube x nbins
-    """
-    y = ws.extractY()[nmonitors:,:]  # exclude monitors - alternatively load with monitors separate?
-    nbins = ws.blocksize()  # 4451
-    y = np.reshape(y, (ntubes, npix_per_tube, nbins))[:,::-1,:]  # ntubes x npix x nbins (note flipped pix along tube)
-    # reverse order of tubes in each bank
-    nbanks = ntubes//ntubes_per_bank
-    for ibank in range(0, nbanks):
-        istart = ibank*ntubes_per_bank
-        iend = (ibank+1)*ntubes_per_bank
-        y[istart:iend,:,:] = y[istart:iend,:,:][::-1,:,:]
-    return y
 
 def save_data(data, filepath):
     """
@@ -68,49 +48,6 @@ def read_data(filepath):
 def save_found_peaks(filepath, indices):
     np.savetxt(filepath, indices, fmt='%d')
 
-# functions to process output of ML peak finding
-
-def findSpectrumIndex(indices, ntubes=1520, npix_per_tube=128, ntubes_per_bank=152, nmonitors=5):
-    """
-    :param indices: indices of found peaks in 3d array
-    :param ntubes: number of tubes in instrument (WISH)
-    :param npix_per_tube: number of detector pixels in each tube
-    :param ntubes_per_bank: number of tubes per bank
-    :param nmonitors: number of monitor spectra (assumed first spectra in ws)
-    :return: list of spectrum indindices of workspace corresponding to indices of found peaks in 3d array
-    """
-    # find bank and then reverse order of tubes in bank
-    ibank = np.floor(indices[:,0]/ntubes_per_bank)
-    itube = ibank*ntubes_per_bank + ((ntubes_per_bank-1) - indices[:,0] % ntubes_per_bank)
-    # flip tube 
-    ipix = (npix_per_tube-1) - indices[:,1]
-    # get spectrum index
-    specIndex = np.ravel_multi_index((itube.astype(int), ipix), 
-        dims=(ntubes, npix_per_tube), order='C') + nmonitors
-    return specIndex.tolist()  # so elements are of type int not numpy.int32
-
-def createPeaksWorkspaceFromIndices(ws, peak_wsname, indices, data):
-    """
-    Create a PeaksWorkspace using indices of peaks found in 3d array. WISH has 4 detectors so put peak in central pixel.
-    Could add peak using avg QLab for peaks at that lambda in all detectors but peak placed in nearest detector anyway
-    for more details see https://github.com/mantidproject/mantid/issues/31944
-    :param ws: MatrixWorkspace of WISH data with xunit wavelength
-    :param peak_wsname: Output name of peaks workspace created
-    :param indices: indices of peaks found in 3d array
-    :param data: 3d array of data
-    :return: PeaksWorkspace
-    """
-    ispec = findSpectrumIndex(indices, *data.shape[0:2])
-    peaks = CreatePeaksWorkspace(InstrumentWorkspace=ws, NumberOfPeaks=0, 
-                OutputWorkspace=peak_wsname)
-    for ipk in range(len(ispec)):
-        wavelength = ws.readX(ispec[ipk])[indices[ipk,2]]
-        # four detectors per spectrum so use one of the central ones
-        detIDs = ws.getSpectrum(ispec[ipk]).getDetectorIDs()
-        idet = (len(detIDs)-1)//2  # pick central pixel
-        AddPeak(PeaksWorkspace=peaks, RunWorkspace=ws, TOF=wavelength, DetectorID=detIDs[idet],
-            BinCount=data[indices[ipk,0], indices[ipk,1], indices[ipk,2]])
-    return peaks
 
 if __name__ == "__main__" or "mantidqt.widgets.codeeditor.execution":
     # 1. reduce data and save 3d numpy array from ML algorithm (bragg-detect)