Update cloud_coverage_pipeline.py

cloud_coverage_pipeline.py

@@ -1,16 +1,16 @@
 # Importing Libraries
-import torch, os
+import os
 import numpy as np
 import cv2
+import torch
 from torch import nn
 import timm
+import pickle
 from transformers import DistilBertModel, DistilBertConfig
 from torch.utils.data import Dataset, DataLoader
 from tqdm.autonotebook import tqdm
-import pickle
 os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
 
-
 # Trained Model Configurations
 CFG = {
     "debug": False,
@@ -40,44 +40,40 @@ CFG = {
     "dropout": 0.1
 }
 
-
 # Loading Finetuned Clip Model to the below class format
 class CLIPModel(nn.Module):
     def __init__(
             self,
-            temperature=CFG["temperature"],
-            image_embedding=CFG["image_embedding"],
-            text_embedding=CFG["text_embedding"],
+            temperature = CFG["temperature"],
+            image_embedding = CFG["image_embedding"],
+            text_embedding = CFG["text_embedding"],
     ):
         super().__init__()
         self.image_encoder = ImageEncoder()
         self.text_encoder = TextEncoder()
-        self.image_projection = ProjectionHead(embedding_dim=image_embedding)
-        self.text_projection = ProjectionHead(embedding_dim=text_embedding)
+        self.image_projection = ProjectionHead(embedding_dim = image_embedding)
+        self.text_projection = ProjectionHead(embedding_dim = text_embedding)
         self.temperature = temperature
 
-
 # Image Encoder Class to extract features using finetuned clip's Resnet Image Encoder
 class ImageEncoder(nn.Module):
-    def __init__(self, model_name=CFG["model_name"], pretrained=CFG["pretrained"], trainable=CFG["trainable"]):
+    def __init__(self, model_name = CFG["model_name"], pretrained = CFG["pretrained"], trainable = CFG["trainable"]):
         super().__init__()
-        self.model = timm.create_model(model_name, pretrained, num_classes=0, global_pool="avg")
+        self.model = timm.create_model(model_name, pretrained, num_classes = 0, global_pool = "avg")
         for p in self.model.parameters():
             p.requires_grad = trainable
 
     def forward(self, x):
         return self.model(x)
 
-
 # Text Encoder - Optional in inference
 class TextEncoder(nn.Module):
-    def __init__(self, model_name=CFG["text_encoder_model"], pretrained=CFG["pretrained"],
-                 trainable=CFG["trainable"]):
+    def __init__(self, model_name = CFG["text_encoder_model"], pretrained = CFG["pretrained"], trainable = CFG["trainable"]):
         super().__init__()
         if pretrained:
             self.model = DistilBertModel.from_pretrained(model_name)
         else:
-            self.model = DistilBertModel(config=DistilBertConfig())
+            self.model = DistilBertModel(config = DistilBertConfig())
 
         for p in self.model.parameters():
             p.requires_grad = trainable
@@ -85,18 +81,17 @@ class TextEncoder(nn.Module):
         self.target_token_idx = 0
 
     def forward(self, input_ids, attention_mask):
-        output = self.model(input_ids=input_ids, attention_mask=attention_mask)
+        output = self.model(input_ids = input_ids, attention_mask = attention_mask)
         last_hidden_state = output.last_hidden_state
         return last_hidden_state[:, self.target_token_idx, :]
 
-
 # Projection Class - Optional in inference
 class ProjectionHead(nn.Module):
     def __init__(
             self,
             embedding_dim,
-            projection_dim=CFG["projection_dim"],
-            dropout=CFG["dropout"]
+            projection_dim = CFG["projection_dim"],
+            dropout = CFG["dropout"]
     ):
         super().__init__()
         self.projection = nn.Linear(embedding_dim, projection_dim)
@@ -114,7 +109,6 @@ class ProjectionHead(nn.Module):
         x = self.layer_norm(x)
         return x
 
-
 # Class to transform image to custom data format
 class SkyImage(Dataset):
     def __init__(self, img, label):
@@ -126,34 +120,34 @@ class SkyImage(Dataset):
 
     def __getitem__(self, idx):
         image = cv2.resize(self.img[idx], (244, 244))
-        # image = cv2.cvtColor(self.img[idx], cv2.COLOR_BGR2RGB)
-        # image = cv2.resize(image, (244, 244))
         image = np.moveaxis(image, -1, 0)
         label = self.img_label[idx]
         return image, label
 
+# Method to initialize CatBoost model
+def initialize_catboost_model():
+    return pickle.load(open("catboost_model.sav", 'rb'))
+
+# Method to initialize CLIP model
+def initialize_clip_model():
+    clip_model = CLIPModel().to(CFG["device"])
+    clip_model.load_state_dict(torch.load("clip_model.pt", map_location = CFG["device"]))
+    clip_model.eval()
+    return clip_model
 
 # Method to extract features from finetuned clip model
 def get_features(clip_model, dataset):
     features, label, embeddings = [], [], []
     with torch.no_grad():
-        for images, labels in tqdm(DataLoader(dataset, batch_size=64)):
+        for images, labels in tqdm(DataLoader(dataset, batch_size = 64)):
             image_input = torch.tensor(np.stack(images)).cpu().float()
             image_features = clip_model.image_encoder(image_input)
             features.append(image_features)
             label.append(labels)
     return torch.cat(features), torch.cat(label).cpu()
 
-
-# Loading Clip and Catboost models
-CTBR_model = pickle.load(open("catboost_model.sav", 'rb'))
-clip_model = CLIPModel().to(CFG["device"])
-clip_model.load_state_dict(torch.load("clip_model.pt", map_location=CFG["device"]))
-clip_model.eval()
-
-
 # Method to calculate cloud coverage
-def predict_cloud_coverage(image):
+def predict_cloud_coverage(image, clip_model, CTBR_model):
     img, lbl = [image], [0]
     # Transforming Data into custom format
     test_image = SkyImage(img, lbl)
@@ -161,4 +155,4 @@ def predict_cloud_coverage(image):
     features, label = get_features(clip_model, test_image)
     # Predicting Cloud Coverage based on extracted features
     pred_cloud_coverage = CTBR_model.predict(features.cpu().numpy())
-    return(round(max(0.0, min(100.0, pred_cloud_coverage[0])), 1))
+    return round(max(0.0, min(100.0, pred_cloud_coverage[0])), 1)