import easyocr
from jiwer import cer, wer
import editdistance
 
reader = easyocr.Reader(['en', 'hi'])
 
def evaluate(image_path, expected_text):
    results = reader.readtext(image_path)
    detected = " ".join([r[1] for r in results]).upper().strip()
    expected = expected_text.upper().strip()
 
    cer_score  = cer(expected, detected)
    wer_score  = wer(expected, detected)
    ed         = editdistance.eval(expected, detected)
    ned        = 1 - (ed / max(len(expected), len(detected)))
    exact      = 1 if expected == detected else 0
 
    return {"CER": cer_score, "WER": wer_score, "NED": ned, "exact": exact}

import requests
import os
import time
import json
from pathlib import Path
 
# Flickr API endpoint
FLICKR_API = "https://api.flickr.com/services/rest/"
 
# Ye saari queries run hongi
QUERIES = [
    "poster design typography",
    "flyer design marketing",
    "wedding invitation card",
    "festival sale banner",
    "product launch poster",
    "Diwali banner design",
    "promotional flyer India",
    "event poster design",
    "banner ad design",
    "graphic design poster",
    # ... aur 24 queries
]
 
def search_flickr(query, api_key, page=1, per_page=100):
    """
    Flickr se images search karo.
    
    params dict mein ye fields important hain:-
    - method: flickr.photos.search — ye Flickr ka search endpoint hai
    - text: hamara query string
    - safe_search: 1 — adult content filter
    - content_type: 1 — sirf photos, illustrations nahi
    - media: photos — sirf images
    - extras: url_l,url_o — large aur original URL bhi do
    - format: json
    """
    params = {
        "method": "flickr.photos.search",
        "api_key": api_key,
        "text": query,
        "safe_search": 1,
        "content_type": 1,
        "media": "photos",
        "extras": "url_l,url_o,views,description",
        "per_page": per_page,
        "page": page,
        "format": "json",
        "nojsoncallback": 1
    }
    
    response = requests.get(FLICKR_API, params=params, timeout=10)
    data = response.json()
    
    # data["photos"]["photo"] mein actual results hain
    return data["photos"]["photo"]
 
def download_image(url, save_path):
    """
    Single image download. Stream=True isliye use karte hain
    kyunki badi images ko memory mein ek saath load nahi karna.
    Chunks mein download hoti hai disk pe.
    """
    if os.path.exists(save_path):
        return True  # already downloaded — skip
    
    try:
        response = requests.get(url, stream=True, timeout=15)
        response.raise_for_status()
        
        with open(save_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        return True
    except Exception as e:
        print(f"Failed: {url} — {e}")
        return False
 
def scrape_flickr(api_key, output_dir, target_per_query=150):
    """
    Main scraping function.
    
    Progress file save karta hai JSON mein — agar script crash ho
    to dobara same images download nahi hogi. Ye resume support hai.
    """
    Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    # Progress tracking
    progress_file = os.path.join(output_dir, "progress.json")
    progress = json.load(open(progress_file)) if os.path.exists(progress_file) else {}
    
    total_downloaded = 0
    
    for query in QUERIES:
        if progress.get(query, 0) >= target_per_query:
            print(f"Skipping '{query}' — already done")
            continue
        
        downloaded_for_query = progress.get(query, 0)
        page = (downloaded_for_query // 100) + 1
        
        print(f"\nQuery: '{query}'")
        
        while downloaded_for_query < target_per_query:
            photos = search_flickr(query, api_key, page=page)
            
            if not photos:
                break
            
            for photo in photos:
                # url_o = original, url_l = large — original prefer karo
                url = photo.get("url_o") or photo.get("url_l")
                if not url:
                    continue
                
                # filename: query_photoId.jpg — unique naam
                safe_query = query.replace(" ", "_")[:30]
                filename = f"{safe_query}_{photo['id']}.jpg"
                save_path = os.path.join(output_dir, filename)
                
                if download_image(url, save_path):
                    downloaded_for_query += 1
                    total_downloaded += 1
                
                if downloaded_for_query >= target_per_query:
                    break
                
                time.sleep(0.5)  # rate limit — Flickr 3600 requests/hour allow karta hai
            
            page += 1
        
        # Progress save karo
        progress[query] = downloaded_for_query
        json.dump(progress, open(progress_file, 'w'))
    
    print(f"\nTotal downloaded: {total_downloaded}")

import torch
from transformers import AutoProcessor, AutoModelForCausalLM
from PIL import Image
import easyocr
import os
import argparse
from pathlib import Path
 
# ─────────────────────────────────────────
# Florence-2 Setup
# ─────────────────────────────────────────
# Florence-2 ek Microsoft ka vision-language model hai.
# Base fine-tuned version (~270M params) use karte hain — fast aur accurate.
# "florence-2-base-ft" mein "ft" = fine-tuned on more tasks.
 
FLORENCE_MODEL = "microsoft/Florence-2-base-ft"
 
def load_florence():
    processor = AutoProcessor.from_pretrained(FLORENCE_MODEL, trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        FLORENCE_MODEL,
        torch_dtype=torch.float32,  # CPU pe float32, GPU pe float16
        trust_remote_code=True
    )
    return model, processor
 
def generate_visual_caption(image_path, model, processor):
    """
    Florence-2 se detailed image description generate karo.
    
    Task "<MORE_DETAILED_CAPTION>" se Florence zyada detail deta hai
    simple "<CAPTION>" se zyada. Ye task token model ko batata hai
    ki kitna detailed respond karna hai.
    """
    image = Image.open(image_path).convert("RGB")
    
    task = "<MORE_DETAILED_CAPTION>"
    
    inputs = processor(
        text=task,
        images=image,
        return_tensors="pt"
    )
    
    with torch.no_grad():
        generated_ids = model.generate(
            input_ids=inputs["input_ids"],
            pixel_values=inputs["pixel_values"],
            max_new_tokens=256,
            num_beams=3,       # beam search — 3 candidates consider karta hai
            do_sample=False    # deterministic output
        )
    
    caption = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
    
    # Florence wraps output in task token — clean karo
    caption = processor.post_process_generation(
        caption, task=task, image_size=(image.width, image.height)
    )
    
    return caption[task].strip()
 
# ─────────────────────────────────────────
# EasyOCR Setup
# ─────────────────────────────────────────
# EasyOCR ek open source OCR library hai.
# ['en', 'hi'] matlab English aur Hindi dono detect karega.
# GPU=False kyunki CPU pe bhi fast enough hai OCR ke liye.
 
def load_ocr():
    return easyocr.Reader(['en', 'hi'], gpu=False)
 
def extract_text_from_image(image_path, reader, min_confidence=0.78):
    """
    Image mein visible text detect karo.
    
    reader.readtext() returns list of tuples:-
    [(bounding_box, text, confidence), ...]
    
    Confidence filter important hai — low confidence detections
    mostly garbage hain (random shapes jo OCR ko text jaisi lagti hain).
    
    Additional filters:-
    - Length filter: 1 se 50 characters ke beech
    - Digit ratio: agar text mostly digits hai to likely price/number
    - Mixed script filter: agar ek word mein Hindi aur English mixed hai
    """
    try:
        results = reader.readtext(image_path)
    except Exception:
        return []
    
    clean_texts = []
    
    for (bbox, text, confidence) in results:
        if confidence < min_confidence:
            continue
        
        text = text.strip()
        
        # Length filter
        if len(text) < 2 or len(text) > 50:
            continue
        
        # Garbage filter — mostly symbols or single characters
        alpha_count = sum(1 for c in text if c.isalpha())
        if alpha_count < 1:
            continue
        
        clean_texts.append(text)
    
    return clean_texts
 
# ─────────────────────────────────────────
# Style Tags from Filename
# ─────────────────────────────────────────
# Filename se design style guess karo.
# Ye rough hai but adds useful context to captions.
 
STYLE_KEYWORDS = {
    "wedding": "wedding design, formal, elegant",
    "diwali": "Diwali festival, Indian festival design",
    "sale": "sale promotion, marketing design",
    "poster": "poster design, large format",
    "flyer": "flyer design, promotional",
    "invitation": "invitation card design",
    "banner": "banner design, horizontal format",
    "product": "product photography, commercial",
}
 
def get_style_tags(filename):
    filename_lower = filename.lower()
    tags = []
    for keyword, tag in STYLE_KEYWORDS.items():
        if keyword in filename_lower:
            tags.append(tag)
    return ", ".join(tags)
 
# ─────────────────────────────────────────
# Final Caption Assembly
# ─────────────────────────────────────────
 
def build_caption(visual_description, ocr_texts, style_tags):
    """
    Teen sources merge karke final training caption banao.
    
    Format:-
    [style tags], [visual description], [text "word1"], [text "word2"]
    
    Example output:-
    "wedding invitation card, formal design, elegant gold border on cream background 
    with floral motifs and serif typography, text 'Save The Date', text 'Mr & Mrs Sharma'"
    """
    parts = []
    
    if style_tags:
        parts.append(style_tags)
    
    if visual_description:
        parts.append(visual_description)
    
    for t in ocr_texts[:6]:  # max 6 text detections per image
        parts.append(f'text "{t}"')
    
    return ", ".join(parts)
 
# ─────────────────────────────────────────
# Main Runner
# ─────────────────────────────────────────
 
def caption_dataset(image_folder, output_folder, skip_done=True):
    Path(output_folder).mkdir(parents=True, exist_ok=True)
    
    print("Loading Florence-2...")
    florence_model, florence_processor = load_florence()
    
    print("Loading EasyOCR...")
    ocr_reader = load_ocr()
    
    image_files = list(Path(image_folder).glob("*.jpg")) + \
                  list(Path(image_folder).glob("*.png"))
    
    print(f"Found {len(image_files)} images")
    
    for i, img_path in enumerate(image_files):
        caption_path = Path(output_folder) / (img_path.stem + ".txt")
        
        if skip_done and caption_path.exists():
            continue
        
        try:
            visual = generate_visual_caption(str(img_path), florence_model, florence_processor)
            ocr = extract_text_from_image(str(img_path), ocr_reader)
            tags = get_style_tags(img_path.name)
            caption = build_caption(visual, ocr, tags)
            
            caption_path.write_text(caption, encoding="utf-8")
            
            if (i + 1) % 50 == 0:
                print(f"Progress: {i+1}/{len(image_files)}")
        
        except Exception as e:
            print(f"Error on {img_path.name}: {e}")

import requests
import os
import json
import time
from PIL import Image
from io import BytesIO
from bs4 import BeautifulSoup
from transformers import AutoProcessor, AutoModelForCausalLM
import easyocr
import torch
from pathlib import Path
import argparse
 
class ScrapeAndCaptionAgent:
    
    def __init__(self, output_dir):
        self.output_dir = Path(output_dir)
        self.images_dir = self.output_dir / "images"
        self.captions_dir = self.output_dir / "captions"
        self.metadata_file = self.output_dir / "metadata.jsonl"
        
        self.images_dir.mkdir(parents=True, exist_ok=True)
        self.captions_dir.mkdir(parents=True, exist_ok=True)
        
        print("Loading Florence-2...")
        self.florence_processor = AutoProcessor.from_pretrained(
            "microsoft/Florence-2-base-ft", trust_remote_code=True
        )
        self.florence_model = AutoModelForCausalLM.from_pretrained(
            "microsoft/Florence-2-base-ft",
            torch_dtype=torch.float32,
            trust_remote_code=True
        )
        
        print("Loading EasyOCR...")
        self.ocr = easyocr.Reader(['en', 'hi'], gpu=False)
    
    # ─────────────────────────────────────────
    # Pinterest Scraping with Context
    # ─────────────────────────────────────────
    
    def scrape_pinterest_with_context(self, query, target=500):
        """
        Pinterest se image + metadata ek saath.
        
        Pin object mein ye fields available hain:-
        - images: different size URLs
        - title: pin ka title
        - description: pin ka description
        - board_name: kis board pe hai
        - alt_text: accessibility text (often very descriptive)
        
        Ye sab milke caption quality kaafi improve karte hain.
        """
        results = []
        bookmark = None
        
        while len(results) < target:
            try:
                pins, bookmark = self._pinterest_search(query, bookmark)
            except Exception as e:
                print(f"Pinterest error: {e}")
                time.sleep(30)
                continue
            
            for pin in pins:
                # Image URL
                images = pin.get("images", {})
                url = (images.get("orig") or images.get("736x") or {}).get("url")
                if not url:
                    continue
                
                # Context metadata — ye caption mein jaayega
                context = {
                    "source": "pinterest",
                    "query": query,
                    "title": pin.get("title", ""),
                    "description": pin.get("description", ""),
                    "alt_text": pin.get("alt_text", ""),
                    "board_name": (pin.get("board") or {}).get("name", ""),
                    "image_url": url
                }
                
                results.append(context)
                time.sleep(1.5)
                
                if len(results) >= target:
                    break
            
            if not bookmark:
                break
        
        return results
    
    # ─────────────────────────────────────────
    # Google Images Scraping
    # ─────────────────────────────────────────
    
    def scrape_google_images(self, query, target=100):
        """
        Google Images se scrape karo via regular search.
        
        Note:- Google images scraping fragile hai. URL structure change hoti rehti hai.
        Isliye ye sirf supplementary source hai, primary nahi.
        
        SerpAPI ya ScraperAPI use karna zyada reliable hoga production mein.
        """
        results = []
        headers = {
            "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"
        }
        
        search_url = f"https://www.google.com/search?q={query}&tbm=isch&hl=en"
        
        try:
            response = requests.get(search_url, headers=headers, timeout=10)
            soup = BeautifulSoup(response.text, 'html.parser')
            
            # Google image tags from search results
            for img in soup.find_all("img", limit=target):
                src = img.get("src", "")
                alt = img.get("alt", "")
                
                if src.startswith("http") and not "google" in src:
                    results.append({
                        "source": "google",
                        "query": query,
                        "image_url": src,
                        "alt_text": alt,
                        "title": "",
                        "description": "",
                        "board_name": ""
                    })
        
        except Exception as e:
            print(f"Google scrape error: {e}")
        
        return results
    
    # ─────────────────────────────────────────
    # Download and Process Single Image
    # ─────────────────────────────────────────
    
    def process_item(self, context_dict):
        """
        Ek image download karo aur uski caption banao.
        
        Steps:-
        1. Image download karo
        2. Size validate karo (minimum 256x256, avoid tiny thumbnails)
        3. Florence-2 se visual description
        4. EasyOCR se text extraction
        5. Context (title, alt, description) se tags
        6. Sab merge karke final caption
        7. Image save karo, caption .txt file save karo
        8. metadata.jsonl mein entry append karo
        """
        url = context_dict["image_url"]
        
        try:
            response = requests.get(url, timeout=15, stream=True)
            response.raise_for_status()
            
            img_bytes = response.content
            image = Image.open(BytesIO(img_bytes)).convert("RGB")
            
            # Minimum size check
            if image.width < 256 or image.height < 256:
                return False
            
        except Exception as e:
            return False
        
        # Unique filename from URL hash
        import hashlib
        url_hash = hashlib.md5(url.encode()).hexdigest()[:12]
        source = context_dict["source"]
        img_filename = f"{source}_{url_hash}.jpg"
        img_path = self.images_dir / img_filename
        
        # Save image
        image.save(str(img_path), "JPEG", quality=95)
        
        # Generate caption parts
        visual_desc = self._florence_caption(image)
        ocr_texts = self._ocr_extract(str(img_path))
        context_text = self._context_to_text(context_dict)
        
        # Merge caption
        parts = []
        if context_text:
            parts.append(context_text)
        if visual_desc:
            parts.append(visual_desc)
        for t in ocr_texts[:5]:
            parts.append(f'text "{t}"')
        
        caption = ", ".join(parts)
        
        # Save caption .txt
        caption_path = self.captions_dir / (img_filename.replace(".jpg", ".txt"))
        caption_path.write_text(caption, encoding="utf-8")
        
        # Append to JSONL
        record = {
            "image": str(img_path),
            "caption": caption,
            "source": source,
            "query": context_dict["query"]
        }
        with open(self.metadata_file, 'a') as f:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")
        
        return True
    
    def _context_to_text(self, ctx):
        """Page context se useful tags banao."""
        parts = []
        if ctx.get("title") and len(ctx["title"]) > 3:
            parts.append(ctx["title"].lower())
        if ctx.get("alt_text") and len(ctx["alt_text"]) > 3:
            parts.append(ctx["alt_text"].lower())
        if ctx.get("board_name") and len(ctx["board_name"]) > 3:
            parts.append(ctx["board_name"].lower())
        return ", ".join(parts[:3])  # max 3 context tags
    
    def _florence_caption(self, pil_image):
        """Florence-2 se visual description."""
        try:
            inputs = self.florence_processor(
                text="<MORE_DETAILED_CAPTION>",
                images=pil_image,
                return_tensors="pt"
            )
            with torch.no_grad():
                ids = self.florence_model.generate(
                    **inputs, max_new_tokens=150, num_beams=3
                )
            result = self.florence_processor.batch_decode(ids, skip_special_tokens=False)[0]
            result = self.florence_processor.post_process_generation(
                result, task="<MORE_DETAILED_CAPTION>",
                image_size=(pil_image.width, pil_image.height)
            )
            return result["<MORE_DETAILED_CAPTION>"].strip()
        except:
            return ""
    
    def _ocr_extract(self, image_path):
        """EasyOCR se text extract karo."""
        try:
            results = self.ocr.readtext(image_path)
            return [r[1].strip() for r in results if r[2] > 0.78 and 2 <= len(r[1]) <= 50]
        except:
            return []