added new model with cosine similarity

app.py

@@ -21,7 +21,7 @@ import albumentations.pytorch as al_pytorch
 import torchvision
 from pl_bolts.models.gans import Pix2Pix
 from pl_bolts.models.gans.pix2pix.components import PatchGAN
-
+import torchvision.models as models
 
 """ Class """
 
@@ -86,6 +86,7 @@ train_16_val_1_plbolts_model_chkpt = (
     "model/lightning_bolts_model/epoch=99-step=89000.ckpt"
 )
 modified_patchgan_chkpt = "model/lightning_bolts_model/modified_patchgan.ckpt"
+
 # model_checkpoint_path = "model/pix2pix_lightning_model/version_0/checkpoints/epoch=199-step=355600.ckpt"
 # model_checkpoint_path = "model/pix2pix_lightning_model/gen.pth"
 
@@ -106,6 +107,112 @@ modified_patchgan_model = PatchGanChanged.load_from_checkpoint(modified_patchgan
 modified_patchgan_model.eval()
 
 
+# Create new class
+class OverpoweredPix2Pix(Pix2Pix):
+    def __init__(self, in_channels, out_channels):
+        super(OverpoweredPix2Pix, self).__init__(
+            in_channels=in_channels, out_channels=out_channels
+        )
+        self._create_inception_score()
+
+    def _gen_step(self, real_images, conditioned_images):
+        # Pix2Pix has adversarial and a reconstruction loss
+        # First calculate the adversarial loss
+        fake_images = self.gen(conditioned_images)
+        disc_logits = self.patch_gan(fake_images, conditioned_images)
+        adversarial_loss = self.adversarial_criterion(
+            disc_logits, torch.ones_like(disc_logits)
+        )
+
+        # calculate reconstruction loss
+        recon_loss = self.recon_criterion(fake_images, real_images)
+        lambda_recon = self.hparams.lambda_recon
+
+        # calculate cosine similarity
+        representations_real = self.feature_extractor(real_images).flatten(1)
+        representations_fake = self.feature_extractor(fake_images).flatten(1)
+        similarity_score_list = self.cosine_similarity(
+            representations_real, representations_fake
+        )
+        cosine_sim = sum(similarity_score_list) / len(similarity_score_list)
+
+        self.log("Gen Cosine Sim Loss ", 1 - cosine_sim.cpu().detach().numpy())
+        # print(adversarial_loss,1-cosine_sim, lambda_recon, recon_loss, )
+
+        return (
+            (adversarial_loss)
+            + (lambda_recon * recon_loss)
+            + (lambda_recon * (1 - cosine_sim))
+        )
+
+    def _create_inception_score(self):
+        # init a pretrained resnet
+        backbone = models.resnet50(pretrained=True)
+        num_filters = backbone.fc.in_features
+        layers = list(backbone.children())[:-1]
+        self.feature_extractor = torch.nn.Sequential(*layers)
+        self.cosine_similarity = torch.nn.CosineSimilarity(dim=1, eps=1e-6)
+
+    def validation_step(self, batch, batch_idx):
+        """Validation step"""
+        real, condition = batch
+        with torch.no_grad():
+            disc_loss = self._disc_step(real, condition)
+            self.log("Valid PatchGAN Loss", disc_loss)
+
+            gan_loss = self._gen_step(real, condition)
+            self.log("Valid Generator Loss", gan_loss)
+
+            #
+            fake_images = self.gen(condition)
+            representations_real = self.feature_extractor(real).flatten(1)
+            representations_fake = self.feature_extractor(fake_images).flatten(1)
+            similarity_score_list = self.cosine_similarity(
+                representations_real, representations_fake
+            )
+            cosine_sim = sum(similarity_score_list) / len(similarity_score_list)
+
+            self.log("Valid Cosine Sim", cosine_sim)
+
+        return {"sketch": condition, "colour": real}
+
+    def validation_epoch_end(
+        self, outputs: Union[EPOCH_OUTPUT, List[EPOCH_OUTPUT]]
+    ) -> None:
+        sketch = outputs[0]["sketch"]
+        colour = outputs[0]["colour"]
+        self.feature_extractor.eval()
+        with torch.no_grad():
+            gen_coloured = self.gen(sketch)
+            representations_gen = self.feature_extractor(gen_coloured).flatten(1)
+            representations_fake = self.feature_extractor(colour).flatten(1)
+
+        similarity_score_list = self.cosine_similarity(
+            representations_gen, representations_fake
+        )
+        similarity_score = sum(similarity_score_list) / len(similarity_score_list)
+
+        grid_image = torchvision.utils.make_grid(
+            [
+                sketch[0],
+                colour[0],
+                gen_coloured[0],
+            ],
+            normalize=True,
+        )
+        self.logger.experiment.add_image(
+            f"Image Grid {str(self.current_epoch)} __ {str(similarity_score)} ",
+            grid_image,
+            self.current_epoch,
+        )
+
+
+cosine_sim_model_chkpt = "model/lightning_bolts_model/cosine_sim_model.ckpt"
+
+cosine_sim_model = OverpoweredPix2Pix.load_from_checkpoint(cosine_sim_model_chkpt)
+cosine_sim_model.eval()
+
+
 def predict(img: Image, type_of_model: str):
     """Create predictions"""
     # transform img
@@ -127,6 +234,8 @@ def predict(img: Image, type_of_model: str):
         model = train_16_val_1_plbolts_model
     elif type_of_model == "train batch size 64, val batch size 16":
         model = train_64_val_16_plbolts_model
+    elif type_of_model == "cosine similarity":
+        model = cosine_sim_model
     else:
         model = modified_patchgan_model
 
@@ -151,6 +260,13 @@ def predict3(img: Image):
     )
 
 
+def predict4(img: Image):
+    return predict(
+        img=img,
+        type_of_model="cosine similarity",
+    )
+
+
 model_input = gr.inputs.Radio(
     [
         "train batch size 16, val batch size 1",
@@ -169,17 +285,19 @@ img_examples = [
     "examples/thesis6.png",
 ]
 
-
 with gr.Blocks() as demo:
     gr.Markdown(" # Colour your sketches!")
     gr.Markdown(" ## Description :")
-    gr.Markdown(" There are three Pix2Pix models in this example:")
+    gr.Markdown(" There are 4 Pix2Pix models in this example:")
     gr.Markdown(" 1. Training batch size is 16 , validation is 1")
     gr.Markdown(" 2. Training batch size is 64 , validation is 16")
     gr.Markdown(
         " 3. PatchGAN is changed, 1 value only instead of 16*16 ;"
         "training batch size is 64 , validation is 16"
     )
+    gr.Markdown(
+        " 4. cosine similarity is also added as a metric in this experiment for the generator. "
+    )
     with gr.Tabs():
         with gr.TabItem("tr_16_val_1"):
             with gr.Row():
@@ -222,6 +340,20 @@ with gr.Blocks() as demo:
                     outputs=image_output3,
                     fn=predict3,
                 )
+        with gr.TabItem("Cosine similarity loss"):
+            with gr.Row():
+                image_input4 = gr.inputs.Image(type="pil")
+                image_output4 = gr.outputs.Image(
+                    type="pil",
+                )
+            colour_4 = gr.Button("Colour it!")
+            with gr.Row():
+                gr.Examples(
+                    examples=img_examples,
+                    inputs=image_input4,
+                    outputs=image_output4,
+                    fn=predict4,
+                )
 
     colour_1.click(
         fn=predict1,
@@ -238,6 +370,11 @@ with gr.Blocks() as demo:
         inputs=image_input3,
         outputs=image_output3,
     )
+    colour_4.click(
+        fn=predict4,
+        inputs=image_input4,
+        outputs=image_output4,
+    )
 
 demo.title = "Colour your sketches!"
 demo.launch()

model/lightning_bolts_model/cosine_sim_model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2987394cad6890877faaf61ade50eada5397c2d1447a48049e8ad3197ea461cc
+size 780630439