Update cloud_coverage_pipeline.py

cloud_coverage_pipeline.py

@@ -11,6 +11,7 @@ from torch.utils.data import Dataset, DataLoader
 from tqdm.autonotebook import tqdm
 os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
 
+
 # Trained Model Configurations
 CFG = {
     "debug": False,
@@ -40,6 +41,7 @@ CFG = {
     "dropout": 0.1
 }
 
+
 # Loading Finetuned Clip Model to the below class format
 class CLIPModel(nn.Module):
     def __init__(
@@ -55,6 +57,7 @@ class CLIPModel(nn.Module):
         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"]):
@@ -66,6 +69,7 @@ class ImageEncoder(nn.Module):
     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"]):
@@ -85,6 +89,7 @@ class TextEncoder(nn.Module):
         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__(
@@ -109,6 +114,7 @@ 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):
@@ -124,16 +130,15 @@ class SkyImage(Dataset):
         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():
+# Method to initialize CatBoost model
+def initialize_models():
+    cbt_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()
-    return clip_model
+    return cbt_model, clip_model
+
 
 # Method to extract features from finetuned clip model
 def get_features(clip_model, dataset):
@@ -146,6 +151,7 @@ def get_features(clip_model, dataset):
             label.append(labels)
     return torch.cat(features), torch.cat(label).cpu()
 
+
 # Method to calculate cloud coverage
 def predict_cloud_coverage(image, clip_model, CTBR_model):
     img, lbl = [image], [0]