!pip install ultralytics transformers pillow torch torchvision
import torch
from ultralytics import YOLO
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
Device
device = "cuda" if torch.cuda.is_available() else "cpu"
Load YOLO (object detection)
yolo = YOLO("yolov8n.pt")
Load BLIP (image captioning)
processor = BlipProcessor.from_pretrained(
"Salesforce/blip-image-captioning-base"
)
blip = BlipForConditionalGeneration.from_pretrained(
"Salesforce/blip-image-captioning-base"
).to(device)
def detect_objects(image_path, conf_threshold=0.6):
results = yolo(image_path)
detected = set()
for r in results:
for box in r.boxes:
conf = float(box.conf[0])
cls = int(box.cls[0])
if conf >= conf_threshold:
label = yolo.names[cls]
detected.add(label)
return list(detected)def generate_caption(image_path):
image = Image.open(image_path).convert("RGB")
inputs = processor(image, return_tensors="pt").to(device)
output = blip.generate(
**inputs,
max_length=30
)
caption = processor.decode(
output[0],
skip_special_tokens=True
)
return captiondef refine_caption(caption, objects):
caption = caption.lower()
# Force object correctness
if "laptop" in objects and "laptop" not in caption:
caption += " using a laptop"
if "person" in objects and not caption.startswith("a person"):
caption = "a person " + caption
return caption.capitalize()def generate_smart_caption(image_path):
objects = detect_objects(image_path)
caption = generate_caption(image_path)
final_caption = refine_caption(caption, objects)
return final_caption, objectsimage_path = "/content/drive/MyDrive/Research/06.png"
caption, objects = generate_smart_caption(image_path)
print("Detected objects:", objects)
print("Final caption:", caption)
