98 lines
4.1 KiB
Python
98 lines
4.1 KiB
Python
|
"""
|
||
|
|
Agricultural Photo Keyword Generator using BLIP-2 model
|
||
|
|
"""
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from transformers import BlipProcessor, BlipForConditionalGeneration
|
||
|
|
from PIL import Image
|
||
|
|
import re
|
||
|
|
from typing import List, Dict, Optional
|
||
|
|
|
||
|
|
class AgricultureKeywordGenerator:
|
||
|
|
def __init__(self):
|
||
|
|
"""Initialize the BLIP-2 model for image captioning and keyword generation"""
|
||
|
|
print("Loading BLIP model for keyword generation...")
|
||
|
|
self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
|
||
|
|
self.model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
|
||
|
|
|
||
|
|
# Agriculture-specific keywords to enhance results
|
||
|
|
self.agriculture_keywords = {
|
||
|
|
'people': ['farmer', 'rancher', 'agricultural worker', 'farm worker', 'dairy farmer'],
|
||
|
|
'animals': ['cow', 'cattle', 'pig', 'chicken', 'livestock', 'dairy cow', 'beef cattle'],
|
||
|
|
'crops': ['corn', 'wheat', 'soybean', 'cotton', 'rice', 'barley', 'oats'],
|
||
|
|
'equipment': ['tractor', 'harvester', 'plow', 'irrigation', 'farm equipment'],
|
||
|
|
'locations': ['field', 'farm', 'barn', 'pasture', 'greenhouse', 'ranch', 'farmland'],
|
||
|
|
'activities': ['planting', 'harvesting', 'milking', 'feeding', 'cultivation']
|
||
|
|
}
|
||
|
|
|
||
|
|
print("Model loaded successfully!")
|
||
|
|
|
||
|
|
def generate_caption(self, image_path: str) -> str:
|
||
|
|
"""Generate a descriptive caption for the image"""
|
||
|
|
try:
|
||
|
|
image = Image.open(image_path).convert('RGB')
|
||
|
|
inputs = self.processor(image, return_tensors="pt")
|
||
|
|
|
||
|
|
with torch.no_grad():
|
||
|
|
out = self.model.generate(**inputs, max_length=50, num_beams=5)
|
||
|
|
|
||
|
|
caption = self.processor.decode(out[0], skip_special_tokens=True)
|
||
|
|
return caption
|
||
|
|
except Exception as e:
|
||
|
|
print(f"Error generating caption for {image_path}: {e}")
|
||
|
|
return ""
|
||
|
|
|
||
|
|
def extract_keywords_from_caption(self, caption: str) -> List[str]:
|
||
|
|
"""Extract agriculture-relevant keywords from caption"""
|
||
|
|
keywords = []
|
||
|
|
caption_lower = caption.lower()
|
||
|
|
|
||
|
|
# Extract keywords from each category
|
||
|
|
for category, terms in self.agriculture_keywords.items():
|
||
|
|
for term in terms:
|
||
|
|
if term in caption_lower:
|
||
|
|
keywords.append(term)
|
||
|
|
|
||
|
|
# Add general descriptive words
|
||
|
|
descriptive_words = re.findall(r'\b(?:green|fresh|organic|rural|outdoor|sunny|large|small|young|old|male|female)\b', caption_lower)
|
||
|
|
keywords.extend(descriptive_words)
|
||
|
|
|
||
|
|
# Remove duplicates and limit to 10 keywords
|
||
|
|
keywords = list(set(keywords))[:10]
|
||
|
|
|
||
|
|
return keywords
|
||
|
|
|
||
|
|
def generate_keywords(self, image_path: str) -> Dict[str, any]:
|
||
|
|
"""Generate keywords and title for an agricultural image"""
|
||
|
|
caption = self.generate_caption(image_path)
|
||
|
|
keywords = self.extract_keywords_from_caption(caption)
|
||
|
|
|
||
|
|
# If we don't have enough keywords, add some generic agricultural terms
|
||
|
|
if len(keywords) < 5:
|
||
|
|
generic_terms = ['agriculture', 'farming', 'rural', 'outdoor', 'field']
|
||
|
|
for term in generic_terms:
|
||
|
|
if term not in keywords:
|
||
|
|
keywords.append(term)
|
||
|
|
if len(keywords) >= 5:
|
||
|
|
break
|
||
|
|
|
||
|
|
return {
|
||
|
|
'caption': caption,
|
||
|
|
'keywords': keywords[:10], # Limit to 10 keywords max
|
||
|
|
'title': self.generate_title(caption)
|
||
|
|
}
|
||
|
|
|
||
|
|
def generate_title(self, caption: str) -> str:
|
||
|
|
"""Generate a product title from the caption"""
|
||
|
|
# Clean up the caption to make it more title-like
|
||
|
|
title = caption.strip()
|
||
|
|
if title and not title[0].isupper():
|
||
|
|
title = title[0].upper() + title[1:]
|
||
|
|
|
||
|
|
# Add "Agricultural" prefix if not agriculture-related
|
||
|
|
agriculture_terms = ['farm', 'agriculture', 'crop', 'livestock', 'rural']
|
||
|
|
if not any(term in title.lower() for term in agriculture_terms):
|
||
|
|
title = f"Agricultural scene: {title}"
|
||
|
|
|
||
|
|
return title
|