inference code added

app.py

@@ -0,0 +1,49 @@
+from flask import Flask, request, jsonify
+from PIL import Image
+import io
+from infer import ImageCaptioningInference
+from models.model import ImageCaptioningModel
+
+app = Flask(__name__)
+
+model_dir = 'model'
+    
+# Initialize inference class
+model = ImageCaptioningModel()
+model.load(model_dir)
+
+inference_model = ImageCaptioningInference(model)
+
+# # Path to the input image
+# image_path = 'test_img.jpg'
+
+# # Perform inference and print the generated caption
+# caption = inference_model.infer_image(image_path)
+# print("Generated Caption:", caption)
+
+@app.route('/')
+def home():
+    return "Welcome to the Flask API"
+
+@app.route('/upload-image', methods=['POST'])
+def upload_image():
+    if 'image' not in request.files:
+        return jsonify({'error': 'No image found in request'})
+    
+    image = request.files['image']
+    # print(image)
+
+    # try :
+    image = Image.open(io.BytesIO(image.read()))
+    # image.show()
+
+    generated_caption = inference_model.infer_image(image)
+
+    return jsonify({'generated_caption': generated_caption})
+
+    
+    # except Exception as e:
+    #     return jsonify({'error': f'{e}'}), 500
+    
+if __name__ == '__main__':
+    app.run(debug=True)

infer.py

@@ -0,0 +1,69 @@
+from PIL import Image
+from models.model import ImageCaptioningModel
+from torchvision import transforms
+import torch
+
+import torch
+from transformers import ViTModel, ViTFeatureExtractor, GPT2LMHeadModel, GPT2Tokenizer
+from PIL import Image
+from config.config import Config
+
+class ImageCaptioningInference:
+    def __init__(self, model):
+        self.model = model
+        self.device = Config.DEVICE
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor()
+        ])
+
+    def infer_image(self, image):
+        # Load and preprocess the image
+        # image = Image.open(image_path)
+        image = self.transform(image).unsqueeze(0).to(self.device)
+
+        # Extract image features
+        image_features = self.model.extract_image_features(image)
+
+        # Generate caption
+        caption = self.generate_caption(image_features)
+        return caption
+    
+    def generate_caption(self, image_features, num_beams=3, max_length=50):
+        # Prepare the image features for input
+        image_features = image_features.unsqueeze(1)  # [batch_size, 1, hidden_size]
+        
+        # Generate caption using beam search
+        output = self.model.gpt2_model.generate(
+            inputs_embeds=image_features,
+            max_length=max_length,
+            num_beams=num_beams,
+            early_stopping=True,
+            pad_token_id=self.model.tokenizer.eos_token_id,
+            bos_token_id=self.model.tokenizer.bos_token_id,
+            eos_token_id=self.model.tokenizer.eos_token_id
+        )
+        
+        # Decode the generated caption
+        caption = self.model.tokenizer.decode(output[0], skip_special_tokens=True)
+        return caption
+
+if __name__ == "__main__":
+    # Path to the saved model directory
+    model_dir = 'model'
+    
+    # Initialize inference class
+    model = ImageCaptioningModel()
+    model.load(model_dir)
+
+    inference_model = ImageCaptioningInference(model)
+    
+    # Path to the input image
+    image_path = 'test_img.jpg'
+
+    image = Image.open(image_path)
+    
+    # Perform inference and print the generated caption
+    caption = inference_model.infer_image(image)
+    print("Generated Caption:", caption)
+

train.py

@@ -7,12 +7,22 @@ from data.dataLoader import ImageCaptionDataset
 from config.config import Config
 from models.model import ImageCaptioningModel
 
-from torchsummary import summary
+import mlflow
+import mlflow.pytorch
 
+# TODO : Implementing Weights and Biases to for project tracking and evaluation and TODO : DVC also for data versioning
 
 
 def train_model(model,dataLoader, optimizer, loss_fn):
 
+    with mlflow.start_run():
+        mlflow.log_params({
+            "epochs": Config.EPOCHS,
+            "batch_size": Config.BATCH_SIZE,
+            "learning_rate": Config.LEARNING_RATE,
+            "device": Config.DEVICE
+        })
+
     model.gpt2_model.train()
     for epoch in range(Config.EPOCHS):
         epoch_loss = 0
@@ -41,10 +51,14 @@ def train_model(model,dataLoader, optimizer, loss_fn):
 
             epoch_loss += loss.item()
 
-        print(f'Epoch {epoch + 1}, Loss: {epoch_loss:.4f}')
+        print(f'Epoch {epoch + 1}, Loss: {epoch_loss/len(dataLoader):.4f}')
+        mlflow.log_metric('loss', epoch_loss/len(dataLoader), step=epoch)
 
     # Save the model    
     model.save('model')
+    # save the artifacts
+    mlflow.log_artifacts('model')
+    mlflow.pytorch.log_model(model.gpt2_model, "models")
 
     # return model
 
@@ -75,4 +89,5 @@ if __name__ == '__main__':
     model = ImageCaptioningModel()
     optimizer = torch.optim.Adam(model.gpt2_model.parameters(), lr=Config.LEARNING_RATE)
     loss_fn = torch.nn.CrossEntropyLoss()
+    mlflow.set_experiment('ImageCaptioning')
     train_model(model, dataloader, optimizer, loss_fn)