Piumi-Udeshika29/caption-generate

!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 caption

def 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, objects

image_path = "/content/drive/MyDrive/Research/06.png"

caption, objects = generate_smart_caption(image_path)

print("Detected objects:", objects) print("Final caption:", caption)