Upload image_caption_generator.py

image_caption_generator.py

@@ -0,0 +1,120 @@
+# -*- coding: utf-8 -*-
+"""Image caption generator.ipynb
+
+Automatically generated by Colaboratory.
+
+Original file is located at
+    https://colab.research.google.com/drive/1kJdblTHuqDn8HCKTuEpoApkN05Gzjpot
+"""
+
+!pip install gradio #used for creating the demo
+
+!pip install timm
+
+!pip install huggingface_hub
+
+from huggingface_hub import notebook_login
+
+notebook_login()
+
+import gradio as gr
+import requests
+import matplotlib.pyplot as plt
+import io
+from PIL import Image
+from matplotlib.patches import Rectangle
+
+#Load model directly
+from transformers import AutoProcessor, BlipForConditionalGeneration, pipeline
+
+# Loading the BLIP model directly which generates the caption
+processor = AutoProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+
+#Using transformers to load DETR model for object detection
+#This model adds a bounding box and label to detected objects
+object_detector = pipeline("object-detection", model="facebook/detr-resnet-50")
+
+#generates the caption for uploaded image
+def caption_generator(input_img):
+  inputs = processor(input_img, return_tensors="pt")
+  out = model.generate(**inputs, max_new_tokens=500)
+  caption = processor.decode(out[0], skip_special_tokens=True)
+  return caption
+
+#function to filter the generated caption checking whether human, cats and/or dogs are present using the labels from the object detection
+#this is the method used in this project
+def filter_caption(object_detection_results):
+  labels = [result['label'] for result in object_detection_results]
+  keywords = ["dog","dogs", "cat","cats", "human","humans","man", "men","woman","women","child","children","adult","adults","person"]
+  return True if any(keyword in labels for keyword in keywords) else False
+
+#function to filter the generated caption checking whether human, cats and/or dogs are present using the generated caption
+#initial method considered
+def filter(caption):
+  #If any of these keywords are present, True is returned
+  keywords = ["dog","dogs", "cat","cats", "human","humans","man", "men","woman","women","child","children","adult","adults","person"]
+  caption = caption.lower()
+  return True if any(keyword in caption for keyword in keywords) else False
+
+#function to create the bounding box and label
+#takes an image and list of results as inputs
+def create_image_bbx_w_label(image, results):
+
+  # Set up the plot
+  fig, ax = plt.subplots(figsize=(12, 8))
+  ax.imshow(image)
+
+  # Plot the bounding boxes and labels
+  for res in results:
+      box = res['box']
+      width = box['xmax'] - box['xmin']
+      height = box['ymax'] - box['ymin']
+
+      rect = Rectangle((box['xmin'], box['ymin']), width, height, linewidth=1, edgecolor='r', facecolor='none')
+      ax.add_patch(rect)
+
+      # Position the label above the rectangle
+      label_position = (box['xmin'], box['ymin'] - 10)
+
+      # Display the label and score
+      label_text = f"{res['label']}: {res['score']:.2f}"
+      ax.text(*label_position, label_text, color='white', fontsize=8, bbox=dict(facecolor='red', alpha=0.5))
+
+  ax.axis('off')
+
+  fname = './img.png'
+  plt.savefig(fname, format='png', bbox_inches='tight', pad_inches=0)
+
+  plt.close(fig)
+  # Load this buffer into a PIL Image
+  pil_img = Image.open(fname)
+
+  # Return the PIL Image object
+  return pil_img
+
+def image_caption_generator(input_image):
+  #detecting objects in image
+  object_detection_results = object_detector(input_image)
+  annotated_img = create_image_bbx_w_label(input_image, object_detection_results)
+
+  #Generating caption of input image
+  caption = caption_generator(input_image)
+  #Filtering the captions for specific case (humans and/or cats/dogs)
+  #filtered_caption = filter(caption) uncomment this if you want to filter using the generated caption
+  filtered_caption = filter_caption(object_detection_results) #uses the generated labels from object detection to filter the captions
+  if filtered_caption:
+    return caption, annotated_img
+  else:
+    return "There are no humans, cats or dogs in this image!", annotated_img
+
+demo = gr.Interface(fn = image_caption_generator,
+                    inputs=[gr.Image(label="Upload image", type="pil")],
+                    outputs=[gr.Textbox(label="Caption"), 'image'],
+                    title="CaptionPlus - Image Caption Generator",
+                    description="Captioning images of humans, cats and/or dogs with object detection",
+                    allow_flagging="never",
+                    examples=["/content/Example.jpg", '/content/OIP.jpg'])
+
+demo.launch(share=True)
+