{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install https://github.com/CellProfiling/HPA-Cell-Segmentation/archive/master.zip\n!pip install dask[complete]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom dask import delayed\nimport os\nimport re\nimport hpacellseg.cellsegmentator as cellsegmentator\nfrom hpacellseg.utils import label_cell, label_nuclei\nimport cv2","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CellProcessor:\n    \n    def __init__(self, img_path, test_data, protein_path=None):\n        if test_data == True:\n            # 3 channels paths for segmentation\n            self.img_path = img_path[:3]\n            # Last channel for protein img\n            self.protein_path = img_path[3][0] \n        else:\n            self.protein_path = protein_path \n            self.masked_img = img_path # Already segmented img \n            \n        self.protein_cells = ()\n        self.masks = () # Needed in test data to evaluate the predictions\n        self.resized_masks = () # Used in crop_cell function\n        self.cropped_cells = ()\n        self.test_data = test_data\n        \n    \n    def add_padding(self, img):\n        img = cv2.copyMakeBorder(img, 224, 224, 224, 224,\n                                 cv2.BORDER_CONSTANT)\n        return img\n    \n    \n    def img_resize(self, img, dim):\n        resized_cell = cv2.resize(img, dim,\n                                  interpolation=cv2.INTER_NEAREST)\n        return resized_cell\n    \n    \n    def crop_cell(self):\n        \n        @delayed\n        def run(protein_cell, bool_mask):\n            # bool_mask is used to create a binary_mask\n            binary_mask = bool_mask.astype(np.uint8)\n            \n            # Make a bounding box around the cell\n            contour,_ = cv2.findContours(\n                binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE\n            )\n            cell_area = max(contour, key=cv2.contourArea)\n            x,y,w,h = cv2.boundingRect(cell_area)\n            cv2.rectangle(protein_cell, (x,y), (x+w, y+h), (0,0), 1)\n            \n            # Cropped the ROI (cell) as square\n            if w > h:\n                h = w # Assign the largest number to both h and w\n            else:\n                w = h\n            cell = protein_cell[y:y+h, x:x+w] # The img size changed\n            \n            # Resize the cell image to 224x224\n            cell = self.img_resize(cell, (224,224))\n            # Normalize to range 0-1\n            #cell = cell.astype('float32')\n            #cell = cell / 255.0\n            \n            return cell\n            \n        \n        for protein_cell, bool_mask in zip(self.protein_cells,\n                                           self.resized_masks):\n            cropped_cell = delayed(run)(protein_cell, bool_mask)\n            self.cropped_cells += (cropped_cell.compute(),)\n            \n    \n    \n    def overlay_protein_image(self, cells_in_img):\n        \n        @delayed(nout=2)\n        def run(cell):\n            # Make a copy of protein_image;\n            # to prevent inplace changes to the image\n            protein_copy = protein_img.copy()\n            # Mask the cell (boolean_mask)\n            mask_cell = cell != 0\n            \n            # Dilate cells to make border around cell\n            kernel = np.ones((3,3))\n            dilated_cell = cv2.dilate(cell, kernel, iterations=1)\n            mask_cell_with_border = dilated_cell != 0\n            \n            # Overlay protein image on cell mask\n            not_cell = ~ mask_cell\n            border = not_cell & mask_cell_with_border\n            protein_copy[not_cell] = 0\n            protein_copy[border] = 255\n            \n            return protein_copy, mask_cell\n        \n    \n        # Read protein image and convert it to grayscale \n        protein_img = cv2.imread(self.protein_path) \n        protein_img = cv2.cvtColor(protein_img, cv2.COLOR_BGR2GRAY)\n        \n        # Resize the protein cell image to 500x500\n        protein_img = self.img_resize(protein_img, (500,500))\n        \n        # Add padding to crop the cells later with same sizes\n        protein_img = self.add_padding(protein_img)\n        \n        # Iterate ovr cells in the image \n        for cell in cells_in_img:\n            protein_copy, mask_cell = delayed(run)(cell)\n            self.protein_cells += (protein_copy.compute(),)\n            self.resized_masks += (mask_cell.compute(),)\n        \n        return self.crop_cell()\n    \n    \n    \n    def cells_separation(self, cell_mask):\n        # Get the pixels of each cell\n        cell_mask = cell_mask.astype(np.uint8)\n        cells_pixels = set(np.ravel(cell_mask))\n        cells_pixels.remove(0)\n        \n        @delayed(nout=2)\n        def run(cell_pixels):\n            # Get the the array of each cell in image\n            cell = np.where(cell_mask == cell_pixels, cell_mask, 0)\n            if self.test_data == True:\n                # Get the mask of each cell\n                mask = cell != 0\n            else:\n                mask = None # Way to decrease memory usage\n            \n            # Resize the cell image to 500x500\n            cell = self.img_resize(cell, (500,500))\n            # Add padding; to crop the cells later with same sizes\n            cell = self.add_padding(cell)\n            \n            return cell, mask\n        \n        \n        # Iterate over cells pixels to separate each cell\n        cells_in_img = ()\n        for cell_pixels in cells_pixels:\n            cell, mask = delayed(run)(cell_pixels)\n            cells_in_img += (cell.compute(),)\n            # Store cell mask for test data\n            if self.test_data == True:\n                self.masks += (mask.compute(),)\n        \n        return self.overlay_protein_image(cells_in_img)\n    \n    \n    \n    def segmentation(self):\n        # HPA-Cell-Segmentation\n        NUC_MODEL = \"./nuclei-model.pth\"\n        CELL_MODEL = \"./cell-model.pth\"\n        segmentator = cellsegmentator.CellSegmentator(\n            NUC_MODEL,\n            CELL_MODEL,\n            scale_factor=0.25,\n            device=\"cuda\",\n            padding=True,\n            multi_channel_model=True\n        )\n        \n        # For nuclei\n        nuc_segmentations = segmentator.pred_nuclei(self.image_paths[2])\n        \n        # For full cells \n        cell_segmentations = segmentator.pred_cells(self.image_paths)\n        \n        # post-processing (masks)\n        nuclei_mask = label_nuclei(nuc_segmentations[0])\n        nuclei_mask, cell_mask = label_cell(\n            nuc_segmentations[0], cell_segmentations[0]\n        )\n        # resize cell_mask img\n        cell_mask = self.img_resize(cell_mask)\n        \n        return self.cells_separation(cell_mask)\n    \n    \n    \n    def run(self):\n        # Training data\n        if self.test_data == False:\n            cell_mask = np.load(self.masked_img)['arr_0']\n            self.cells_separation(cell_mask)\n            return self.cropped_cells\n        # Test data\n        else:\n            self.segmentation()\n            return (self.masks, self.cropped_cells)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class DFConstructor:\n    \n    def __init__(self, imgs_paths, imgs_ids=None,\n                 imgs_labels=None, protein_imgs_paths=None,\n                 test_data=True):\n        \n        self.paths = imgs_paths\n        self.IDs = imgs_ids\n        self.labels = imgs_labels\n        self.protein_paths = protein_imgs_paths\n        self.test_data = test_data\n        \n        \n    def construct(self):\n        batch = 500\n        for idx in range(0, len(self.paths), batch):\n            \n            # For test data\n            if self.test_data == True:\n                for path, ID in zip(self.paths[idx:idx+batch],\n                                   self.IDs[idx:idx+batch]):\n                    # Get protein_cells, masks, cropped_cells\n                    processing = CellProcessor(img_path=path,\n                                               test_data=self.test_data)\n                    masks, cropped_cells = processing.run()\n                    # Create dataframe for cells info\n                    cells_num = len(cropped_cells) # number of cells in one image\n                    df = pd.DataFrame({'ID': [ID] * cells_num,\n                                       'mask': masks,\n                                       'cells': cropped_cells})\n                    yield df\n                    \n            # For train data\n            else:\n                for path, label, protein_path, ID in zip(self.paths[idx:idx+batch],\n                                                     self.labels[idx:idx+batch],\n                                                     self.protein_paths[idx:idx+batch],\n                                                     self.IDs[idx:idx+batch]):\n                    \n                    processing = CellProcessor(path, self.test_data,\n                                               protein_path)\n                    cropped_cells = processing.run()\n                    \n                    # Create dataframe for cells info\n                    cells_num = len(cropped_cells)\n                    df = pd.DataFrame({'label': [label] * cells_num,\n                                       'ID': [ID] * cells_num,\n                                       'cells': cropped_cells})\n                    yield df\n                    \n    \n    def start(self):\n        # Get the whole dataframe of all images\n        whole_df = pd.concat(list(self.construct()), ignore_index=True)\n        whole_df.reset_index(drop=True, inplace=True)\n            \n        return whole_df    ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load train info df (image id, label): images are 8-bit\ntrain_info_path = '../input/hpa-single-cell-image-classification/train.csv'\ntrain_info = pd.read_csv(train_info_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Iterate over hpa masked images to get the ids and path\nhpa_mask_path = '../input/hpa-mask/hpa_cell_mask'\ndata = []\nfor hpa_mask_dir, _, masked_images in os.walk(hpa_mask_path):\n    for masked_image in masked_images:\n        pattern = r'(.*)(.npz)'\n        id = re.match(pattern, masked_image).group(1)\n        masked_image_path = os.path.join(hpa_mask_dir,masked_image)\n        data.append((id, masked_image_path))\n\n# Create hpa_mask_df (ID, masked_path, label)\ndf = pd.DataFrame(data, columns=['ID', 'path'])\nhpa_mask_df = train_info.merge(df, on='ID')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Get the protein images path from train_info df\ntrain_data_path = '../input/hpa-single-cell-image-classification/train/'\nimages_ids = hpa_mask_df.ID.tolist()\nprotein_images_paths = [train_data_path + id + '_green.png' for id in images_ids]","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\npaths = hpa_mask_df.path.tolist()\nlabels = hpa_mask_df.Label.tolist()\nconstruct = DFConstructor(imgs_paths=paths[:1], imgs_ids=images_ids[:1],\n                          imgs_labels=labels[:1],\n                          protein_imgs_paths=protein_images_paths[:1],\n                          test_data=False)\ntrain_df = construct.start()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir hpa_data_results ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from PIL import Image\nfor idx, row in train_df.iterrows():\n    img = Image.fromarray(row.cells)\n    img.save(f'./hpa_data_results/{row.ID}_{idx}.png')","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#!zip -r hpa_data_results.zip ./hpa_data_results\n!tar -zcvf hpa_data_results.tar.gz ./hpa_data_results","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}