fix #4 : add gpu support (#34)

Author: rtharungowda

devolearn/embryo_generator_model/embryo_generator_model.py

@@ -65,7 +65,7 @@ def forward(self, input):
 
 
 class embryo_generator_model():   
-    def __init__(self, mode = "cpu"):
+    def __init__(self, device = "cpu"):
 
         """
         ngf = size of feature maps in generator
@@ -74,11 +74,11 @@ def __init__(self, mode = "cpu"):
         Do not tweak these unless you're changing the Generator() with a new model with a different architecture. 
     
         """
-
+        self.device = device
         self.ngf = 128 ## generated image size 
         self.nz = 128
         self.nc = 1
-        self.generator= Generator(self.ngf, self.nz, self.nc)
+        self.generator = Generator(self.ngf, self.nz, self.nc)
         self.model_url = "https://raw.githubusercontent.com/DevoLearn/devolearn/master/devolearn/embryo_generator_model/embryo_generator.pth"
         self.model_name = "embryo_generator.pth"
         self.model_dir = os.path.dirname(__file__)
@@ -87,15 +87,16 @@ def __init__(self, mode = "cpu"):
 
         try:
             # print("model already downloaded, loading model...")
-            self.generator.load_state_dict(torch.load(self.model_dir + "/" + self.model_name, map_location= "cpu"))
+            self.generator.load_state_dict(torch.load(self.model_dir + "/" + self.model_name, map_location= self.device))
         except:
             print("model not found, downloading from: ", self.model_url)
             if os.path.isdir(self.model_dir) == False:
                 os.mkdir(self.model_dir)
             filename = wget.download(self.model_url, out= self.model_dir)
             # print(filename)
-            self.generator.load_state_dict(torch.load(self.model_dir + "/" + self.model_name, map_location= "cpu"))
-
+            self.generator.load_state_dict(torch.load(self.model_dir + "/" + self.model_name, map_location= self.device))
+        
+        self.generator.to(self.device)
 
 
 
@@ -118,7 +119,7 @@ def generate(self, image_size = (700,500)):
         generated image to the desired size. 
         """
         with torch.no_grad():
-            noise = torch.randn([1,128,1,1])
+            noise = torch.randn([1,128,1,1]).to(self.device)
             im = self.generator(noise)[0][0].cpu().detach().numpy()
         im = cv2.resize(im, image_size)
         im = 255 - cv2.convertScaleAbs(im, alpha=(255.0))   ## temporary fix against inverted images 

devolearn/embryo_segmentor/embryo_segmentor.py

@@ -63,19 +63,18 @@ def generate_centroid_image(thresh):
     return centroid_image, centroids
 
 class embryo_segmentor():
-    def __init__(self):
+    def __init__(self, device = "cpu"):
 
         """
         Segments the c. elegans embryo from images/videos, 
         depends on segmentation-models-pytorch for the model backbone
 
         """
-
+        self.device = device
         self.ENCODER = 'resnet18'
         self.ENCODER_WEIGHTS = 'imagenet'
         self.CLASSES = ["nucleus"]
         self.ACTIVATION = 'sigmoid'
-        self.DEVICE = 'cpu'
         self.in_channels = 1
         self.model_url = "https://github.com/DevoLearn/devolearn/raw/master/devolearn/embryo_segmentor/3d_segmentation_model.pth"
         self.model_name = "3d_segmentation_model.pth"
@@ -92,16 +91,16 @@ def __init__(self):
 
         try:
             # print("model already downloaded, loading model...")
-            self.model = torch.load(self.model_dir + "/" + self.model_name, map_location= "cpu") 
+            self.model = torch.load(self.model_dir + "/" + self.model_name, map_location= self.device) 
         except:
             print("model not found, downloading from:", self.model_url)
             if os.path.isdir(self.model_dir) == False:
                 os.mkdir(self.model_dir)
             filename = wget.download(self.model_url, out= self.model_dir)
             # print(filename)
-            self.model = torch.load(self.model_dir + "/" + self.model_name, map_location= "cpu") 
-
+            self.model = torch.load(self.model_dir + "/" + self.model_name, map_location= self.device) 
 
+        self.model.to(self.device)
         self.model.eval()
 
         self.mini_transform = transforms.Compose([
@@ -129,7 +128,7 @@ def predict(self, image_path, pred_size = (350,250), centroid_mode = False):
         """
 
         im = cv2.imread(image_path,0)
-        tensor = self.mini_transform(im).unsqueeze(0)
+        tensor = self.mini_transform(im).unsqueeze(0).to(self.device)
         res = self.model(tensor).detach().cpu().numpy()[0][0]
         res = cv2.resize(res,pred_size)
         if centroid_mode == False:
@@ -165,7 +164,7 @@ def predict_from_video(self, video_path, pred_size = (350,250), save_folder = "p
         if notebook_mode == True:
             for i in tqdm_notebook(range(len(images)), desc = "saving predictions: "):    
                 save_name = save_folder + "/" + str(i) + ".jpg"
-                tensor = self.mini_transform(images[i]).unsqueeze(0)
+                tensor = self.mini_transform(images[i]).unsqueeze(0).to(self.device)
                 res = self.model(tensor).detach().cpu().numpy()[0][0]
 
                 if centroid_mode == True:
@@ -177,7 +176,7 @@ def predict_from_video(self, video_path, pred_size = (350,250), save_folder = "p
         else :
             for i in tqdm(range(len(images)), desc = "saving predictions: "):
                 save_name = save_folder + "/" + str(i) + ".jpg"
-                tensor = self.mini_transform(images[i]).unsqueeze(0)
+                tensor = self.mini_transform(images[i]).unsqueeze(0).to(self.device)
                 res = self.model(tensor).detach().cpu().numpy()[0][0]
 
                 if centroid_mode == True:

devolearn/lineage_population_model/lineage_population_model.py

@@ -25,8 +25,9 @@
 """
 
 class lineage_population_model():   
-    def __init__(self, mode = "cpu"):
-        self.mode = mode
+    def __init__(self, device = "cpu"):
+        
+        self.device = device
         self.model = models.resnet18(pretrained = True)
         self.model.fc = nn.Linear(512, 7)  ## resize last layer
         self.model_dir = os.path.dirname(__file__)
@@ -39,14 +40,14 @@ def __init__(self, mode = "cpu"):
 
         try:
             # print("model already downloaded, loading model...")
-            self.model.load_state_dict(torch.load(self.model_dir + "/" + self.model_name, map_location= "cpu"))
+            self.model.load_state_dict(torch.load(self.model_dir + "/" + self.model_name, map_location= self.device))
         except:
             print("model not found, downloading from:", self.model_url)
             filename = wget.download(self.model_url, out= self.model_dir)
             # print(filename)
-            self.model.load_state_dict(torch.load(self.model_dir + "/" + self.model_name, map_location= "cpu"))
-
+            self.model.load_state_dict(torch.load(self.model_dir + "/" + self.model_name, map_location= self.device))
 
+        self.model.to(self.device)
         self.model.eval()
 
         self.transforms = transforms.Compose([
@@ -78,8 +79,7 @@ def predict(self, image_path):
 
         image = cv2.imread(image_path, 0)
         image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
-        tensor = self.transforms(image).unsqueeze(0)
-        
+        tensor = self.transforms(image).unsqueeze(0).to(self.device)
         pred = self.model(tensor).detach().cpu().numpy().reshape(1,-1)
 
         pred_scaled = (self.scaler.inverse_transform(pred).flatten()).astype(np.uint8)
@@ -144,13 +144,13 @@ def predict_from_video(self, video_path, csv_name  = "foo.csv", save_csv = False
 
         if notebook_mode == True:
             for i in tqdm_notebook(range(len(images)), desc='Predicting from video file:  :'):
-                tensor = self.transforms(images[i]).unsqueeze(0)
+                tensor = self.transforms(images[i]).unsqueeze(0).to(self.device)
                 pred = self.model(tensor).detach().cpu().numpy().reshape(1,-1)
                 pred_scaled = (self.scaler.inverse_transform(pred).flatten()).astype(np.uint8)
                 preds.append(pred_scaled)
         else :
             for i in tqdm(range(len(images)), desc='Predicting from video file:  :'):
-                tensor = self.transforms(images[i]).unsqueeze(0)
+                tensor = self.transforms(images[i]).unsqueeze(0).to(self.device)
                 pred = self.model(tensor).detach().cpu().numpy().reshape(1,-1)
                 pred_scaled = (self.scaler.inverse_transform(pred).flatten()).astype(np.uint8)
                 preds.append(pred_scaled)
@@ -180,7 +180,7 @@ def predict_from_video(self, video_path, csv_name  = "foo.csv", save_csv = False
 
 
 
-    def create_population_plot_from_video(self, video_path, save_plot = False, plot_name = "plot.png", ignore_first_n_frames = 0, ignore_last_n_frames = 0 ):
+    def create_population_plot_from_video(self, video_path, save_plot = False, plot_name = "plot.png", ignore_first_n_frames = 0, ignore_last_n_frames = 0, notebook_mode = False):
 
         """
         inputs{
@@ -189,6 +189,7 @@ def create_population_plot_from_video(self, video_path, save_plot = False, plot_
             plot_name <str> = filename of the plot image to be saved 
             ignore_first_n_frames <int> = number of frames to drop in the start of the video 
             ignore_last_n_frames <int> = number of frames to drop in the end of the video 
+            notebook_mode <bool> = toogle between script(False) and notebook(True), for better user interface
         }
 
         outputs{
@@ -198,7 +199,7 @@ def create_population_plot_from_video(self, video_path, save_plot = False, plot_
         plots all the predictions from a video into a matplotlib.pyplot 
         
         """
-        df = self.predict_from_video(video_path, ignore_first_n_frames = ignore_first_n_frames, ignore_last_n_frames = ignore_last_n_frames )  
+        df = self.predict_from_video(video_path, ignore_first_n_frames = ignore_first_n_frames, ignore_last_n_frames = ignore_last_n_frames, notebook_mode = notebook_mode)  
 
         labels = ["A", "E", "M", "P", "C", "D", "Z"]
 

devolearn/tests/test.py

@@ -20,7 +20,7 @@ class test(unittest.TestCase):
     def test_lineage_population_model(self):
         test_dir = os.path.dirname(__file__)
 
-        model = lineage_population_model(mode = "cpu")
+        model = lineage_population_model(device = "cpu")
         pred = model.predict(image_path = test_dir + "/" + "sample_data/images/embryo_sample.png")
         self.assertTrue(isinstance(pred, dict), "should be dict")