Fourth Commit

Author: SIDDHARTA

.ipynb_checkpoints/Exploratory_Data_Analysis-checkpoint.ipynb

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7ab6fafee16ec4daca6cf62a57314ee6886165987afee8e04d7f9f559617a3f6
-size 126339797
+oid sha256:9e76358176631bed1f29a87bdf5d1b9c0a7da1d25c461b69c0e6c88e0c100ed5
+size 126373167

.ipynb_checkpoints/Inference_Script-checkpoint.ipynb

@@ -46,8 +46,8 @@ def __getitem__(self, index):
 
     def make_mask(self,row_id):
         name = self.df.iloc[row_id].Name
-        annot = "\\".join([self.annot_folder, name])
-        fname = "\\".join([self.image_folder, (name[:-4]+'.jpg')])
+        annot = "\\".join([self.annot_folder, name+'.xml'])
+        fname = "\\".join([self.image_folder, (name+'.jpg')])
         labels = self.df.iloc[row_id][1:7].to_dict()
 
         tree = ET.parse(annot)

.ipynb_checkpoints/Inference_Script_Real_Time-checkpoint.ipynb

@@ -0,0 +1,75 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[2]:
+
+
+import xml.etree.ElementTree as ET
+import os
+import cv2
+import pandas as pd
+import numpy as np
+from itertools import product
+
+
+# In[ ]:
+
+def create_filepaths(path):
+    df = pd.DataFrame()
+    for (dirpath, dirnames,filenames) in os.walk(path):
+        for filename in filenames:
+            temp_path = "\\".join([path, filename])
+            tree = ET.parse(temp_path)
+            root = tree.getroot()
+            dict1 = dict()
+            ls = []
+            for description in root.iter('name'):
+                ls.append(description.text)
+            res = np.array(ls)
+            res = np.unique(res)
+            ls = res.tolist()
+            dict1['Name'] = filename[:-4]
+            for ele in ls:
+                dict1[ele] = 1
+            df = df.append(dict1, ignore_index= True)
+    df = df.replace(np.nan, 0)
+    df = df[['Name',  'crazing', 'patches', 'inclusion', 'pitted_surface',  'rolled-in_scale', 'scratches']]
+    df['Number_of_Defects'] = df.drop('Name',axis=1).sum(axis=1) 
+        
+    return df  
+
+def make_mask(annot, labels):
+    tree = ET.parse(annot)
+    root = tree.getroot()
+    #extract image  dimensions
+    width = int(root.find('.//size/width').text)
+    height = int(root.find('.//size/height').text)
+    masks = np.zeros((height, width, len(labels)), dtype=np.uint8) # float32 is V.Imp
+        
+    for idx, label in enumerate(labels):
+        if labels.get(label) == 1:
+            boxes = []
+            for obj in root.findall('.//object'):
+                if obj.find('name').text == label:                        
+                    for box in obj.findall('.//bndbox'): 
+                        xmin = int(box.find('xmin').text)
+                        ymin = int(box.find('ymin').text)
+                        xmax = int(box.find('xmax').text)
+                        ymax = int(box.find('ymax').text)
+                        coors = [xmin, ymin, xmax, ymax]
+                        boxes.append(coors)    
+                            
+            mask = np.zeros((height, width), dtype=np.uint8)
+            for i in range(len(boxes)):
+                box = boxes[i]
+                row_s, row_e = box[1], box[3]
+                col_s, col_e = box[0], box[2]
+                row_corr = [*range(row_s, row_e)]
+                col_corr = [*range(col_s, col_e)]
+                coords = np.array(list(product(row_corr, col_corr)))
+                mask[coords[:,0], coords[:,1]] = 1
+                    
+            masks[:, :, idx] = mask.reshape(height, width, order='F')
+         
+    return masks
+

.ipynb_checkpoints/Train_Segmentation_Model-checkpoint.ipynb

@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+# In[2]:
+
+
+from Resnet_UNet import ResNetUNet
+import torch.nn as nn
+from torchvision import models
+import torch
+from collections import OrderedDict
+import argparse
+
+
+# In[2]:
+
+
+def main():
+    my_parser = argparse.ArgumentParser(description='Convert pytorch model to ONNX format.')
+    
+    my_parser.add_argument('Model_Path',
+                       metavar='model_path',
+                       type=str,
+                       help='The path to the saved pytorch model.')
+    
+    args = my_parser.parse_args()
+    
+    ckpt_path = args.Model_Path
+    model = ResNetUNet(n_class=6, base_model=models.resnet18(pretrained=True))
+    state = torch.load(ckpt_path, map_location=lambda storage, loc: storage)
+    new_state_dict = OrderedDict()
+    for k, v in state["state_dict"].items():
+        name = k[7:]
+        new_state_dict[name] = v
+    model.load_state_dict(new_state_dict)
+    model.eval()
+
+    input_var = torch.rand(1, 3, 128, 128)  # Use half of the original resolution.
+    batch_size = 5
+    # Export the model
+    torch.onnx.export(model,                 # model being run
+                  input_var,                 # model input (or a tuple for multiple inputs)
+                  "./ONNX_models/Unet.onnx", # where to save the model (can be a file or file-like object)
+                  export_params=True,        # store the trained parameter weights inside the model file
+                  opset_version=11,          # the ONNX version to export the model to
+                  do_constant_folding=True,  # whether to execute constant folding for optimization
+                  input_names = ['input'],   # the model's input names
+                  output_names = ['output']) # the model's output names
+    
+    print("Successfully converted the model to onnx format.")
+
+
+# In[ ]:
+
+
+if __name__ == "__main__":
+    main()
+

Exploratory_Data_Analysis.ipynb

@@ -6,8 +6,6 @@
 
 import pandas as pd
 import numpy as np
-from Utilities.Data_Retriever_Seg import DataRetriever
-from Utilities.Meter import Meter
 import torch
 import torch.nn as nn
 from sklearn.model_selection import train_test_split
@@ -19,7 +17,9 @@
 import gc
 from tqdm import tqdm
 from os import path, walk
-import xml.etree.ElementTree as ET
+from Utilities.Extract_masks import create_filepaths
+from Utilities.Data_Retriever_Seg import DataRetriever
+from Utilities.Meter import Meter
 
 
 # In[2]:
@@ -145,8 +145,8 @@ def provider(self,
         '''Returns dataloader for the model training'''
         image_folder = '.\IMAGES'
         annot_folder = '.\ANNOTATIONS'
-        df = self.create_filepaths(annot_folder)
-        train_df, val_df = train_test_split(df, test_size=0.3, stratify=df["defects"], random_state=69)
+        df = create_filepaths(annot_folder)
+        train_df, val_df = train_test_split(df, test_size=0.3, stratify=df["Number_of_Defects"], random_state=69)
         df = train_df if phase == "train" else val_df
         image_dataset = DataRetriever(df, image_folder, annot_folder, mean, std, phase)
         dataloader = DataLoader(
@@ -159,26 +159,3 @@ def provider(self,
 
         return dataloader
 
-    def create_filepaths(self, path):
-        df = pd.DataFrame()
-        for (dirpath, dirnames,filenames) in walk(path):
-            for filename in filenames:
-                temp_path = "\\".join([path, filename])
-                tree = ET.parse(temp_path)
-                root = tree.getroot()
-                dict1 = dict()
-                ls = []
-                for description in root.iter('name'):
-                    ls.append(description.text)
-                res = np.array(ls)
-                res = np.unique(res)
-                ls = res.tolist()
-                dict1['Name'] = filename
-                for ele in ls:
-                    dict1[ele] = 1
-                df = df.append(dict1, ignore_index= True)
-        df = df.replace(np.nan, 0)
-        df['defects'] = df.drop('Name',axis=1).sum(axis=1) 
-        
-        return df   
-