#!/usr/bin/env python3 """ Stage 2: Gemini Vision Classification Classifies images using Google Gemini with 5 classification tasks """ import os import json from PIL import Image from io import BytesIO import concurrent.futures from pathlib import Path import time import logging from typing import Dict, Any import mimetypes import random # Gemini SDK from google import genai from google.genai.errors import ServerError from google.genai.types import ( Blob, Part, Content, GenerateContentConfig, ) # Set up logging logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) GEMINI_API_KEY_FALLBACK = "AIzaSyBCgkB2nRaRNgbl06MBu1I_xHiuXSUQMHA" def check_api_key(): """Ensure Google API key is set for Gemini client.""" if not os.getenv('GOOGLE_API_KEY'): # Use provided key as fallback if env not set os.environ['GOOGLE_API_KEY'] = GEMINI_API_KEY_FALLBACK return True def _guess_mime_type(image_path: str) -> str: guessed, _ = mimetypes.guess_type(image_path) if guessed: return guessed try: with Image.open(image_path) as im: fmt = (im.format or '').lower() if fmt in ('jpeg', 'jpg'): return 'image/jpeg' if fmt == 'png': return 'image/png' if fmt == 'webp': return 'image/webp' if fmt == 'gif': return 'image/gif' except Exception: pass return 'application/octet-stream' def _gemini_call_with_retry(contents, cfg, max_attempts: int = 5): """Call Gemini with retries on server/errors.""" api_key = os.getenv('GOOGLE_API_KEY') or GEMINI_API_KEY_FALLBACK for attempt in range(max_attempts): try: client = genai.Client(api_key=api_key) return client.models.generate_content( model="models/gemini-2.5-flash", contents=contents, config=cfg, ) except ServerError as e: sleep_s = (2 ** attempt) + random.random() logger.warning(f"Gemini server error attempt {attempt+1}/{max_attempts}: {e}; retrying in {sleep_s:.1f}s") time.sleep(sleep_s) except Exception as e: sleep_s = (2 ** attempt) + random.random() logger.warning(f"Gemini error attempt {attempt+1}/{max_attempts}: {e}; retrying in {sleep_s:.1f}s") time.sleep(sleep_s) raise RuntimeError(f"Persistent Gemini error after {max_attempts} tries") def classify_image_with_gemini(image_path: str, caption: str, max_retries: int = 3) -> Dict[str, Any]: """Use Google Gemini to classify an image with structured JSON output.""" prompt = f""" Analyze this image with caption: "{caption}" Please answer the following 5 classification questions and respond ONLY with valid JSON: 1. Overall Description. 2. Is the image product display / low quality advertisement / e-commerce product? Answer: "yes" or "no" 3. Is the image computer screenshot with many text overlays? Answer: "yes" or "no" 4. In what category is the image? Choose one from: "animals", "artifacts", "people", "outdoor_scenes", "illustrations", "vehicles", "food_and_beverage", "arts", "abstract", "produce_and_plants", "indoor_scenes" 5. Would you say the image is interesting? Answer: "yes" or "no" 6. Do you think the photo/image was made by a professional photographer? Answer: "yes" or "no" IMPORTANT: Respond with ONLY a valid JSON object with these exact keys. Do not include any other text or explanation: {{ "overall_description": "...", "is_product_advertisement": "yes", "is_screenshot_with_text": "no", "category": "animals", "is_interesting": "no", "is_professional": "yes" }} """ default_response = { "overall_description": "...", "is_product_advertisement": "...", "is_screenshot_with_text": "...", "category": "...", "is_interesting": "...", "is_professional": "..." } try: with open(image_path, 'rb') as f: image_bytes = f.read() mime_type = _guess_mime_type(image_path) image_blob = Blob(mime_type=mime_type, data=image_bytes) user_content = Content( role="user", parts=[ Part(text=prompt), Part(inline_data=image_blob), ], ) contents = [user_content] cfg = GenerateContentConfig(max_output_tokens=2500, temperature=0) resp = _gemini_call_with_retry(contents, cfg, max_attempts=max_retries) logger.debug(f"Gemini response type: {type(resp)}") # Detailed debugging of response structure logger.debug(f"Response.text: {getattr(resp, 'text', 'NO_TEXT_ATTR')}") logger.debug(f"Response.candidates: {getattr(resp, 'candidates', 'NO_CANDIDATES_ATTR')}") if hasattr(resp, 'candidates') and resp.candidates: logger.debug(f"Number of candidates: {len(resp.candidates)}") for i, candidate in enumerate(resp.candidates): logger.debug(f"Candidate {i}: {candidate}") if hasattr(candidate, 'content'): logger.debug(f"Candidate {i} content: {candidate.content}") if hasattr(candidate.content, 'parts'): logger.debug(f"Candidate {i} parts: {candidate.content.parts}") # Check for prompt_feedback which might indicate filtering if hasattr(resp, 'prompt_feedback'): logger.debug(f"Prompt feedback: {resp.prompt_feedback}") # Extract text from Gemini response content_text = "" try: # Try the .text property first if hasattr(resp, 'text') and resp.text: content_text = resp.text logger.debug(f"Got text from .text property: {content_text[:100]}...") else: # Fallback: extract from candidates if resp.candidates and len(resp.candidates) > 0: candidate = resp.candidates[0] if hasattr(candidate, 'content') and candidate.content: if hasattr(candidate.content, 'parts') and candidate.content.parts: for part in candidate.content.parts: if hasattr(part, 'text') and part.text: content_text += part.text logger.debug(f"Got text from candidate part: {part.text[:100]}...") except Exception as e: logger.error(f"Error extracting text from Gemini response: {e}") raise e if not content_text: logger.error(f"Empty response from Gemini") return default_response content_text = content_text.strip() start_idx = content_text.find('{') end_idx = content_text.rfind('}') + 1 if start_idx == -1 or end_idx == 0: logger.error(f"No JSON found in response: {content_text}") return default_response json_content = content_text[start_idx:end_idx] classification = json.loads(json_content) required_keys = [ "overall_description", "is_product_advertisement", "is_screenshot_with_text", "category", "is_interesting", "is_professional", ] missing_keys = [key for key in required_keys if key not in classification] if missing_keys: logger.warning(f"Missing keys in classification: {missing_keys}") for key in missing_keys: classification[key] = default_response[key] return classification except json.JSONDecodeError as e: logger.error(f"JSON parsing error: {e}") return default_response except Exception as e: logger.error(f"Gemini classification error: {e}") return default_response def classify_single_image(metadata_file: Path) -> bool: """Classify a single image and save results""" try: # Load metadata with open(metadata_file, 'r', encoding='utf-8') as f: metadata = json.load(f) idx = metadata['idx'] image_path = metadata['image_path'] caption = metadata['caption'] # Check if image exists if not os.path.exists(image_path): logger.error(f"Image not found: {image_path}") return False # Classify with Gemini classification = classify_image_with_gemini(image_path, caption) # Add classification to metadata metadata['classification'] = classification metadata['stage2_complete'] = True # Save updated metadata new_metadata_file = metadata_file.with_name(f'meta_{idx}_stage2.json') with open(new_metadata_file, 'w', encoding='utf-8') as f: json.dump(metadata, f, indent=2, ensure_ascii=False) logger.info(f"Classified image {idx}") return True except Exception as e: logger.error(f"Error classifying {metadata_file}: {e}") return False def classify_all_images(max_workers: int = 2): """Classify all downloaded images with parallel processing""" logger.info("Starting image classification...") # Get all metadata files metadata_dir = Path('./data/metadata') metadata_files = list(metadata_dir.glob('meta_*.json')) if not metadata_files: logger.error("No metadata files found. Run stage 1 first.") return successful = 0 with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: futures = [executor.submit(classify_single_image, meta_file) for meta_file in metadata_files] for future in concurrent.futures.as_completed(futures): if future.result(): successful += 1 # Rate limiting for API calls time.sleep(1.0) # 1 second between API calls to avoid rate limits logger.info(f"Successfully classified {successful}/{len(metadata_files)} images") # Save summary summary = { "total_images": len(metadata_files), "successful_classifications": successful, "stage": "classification_complete" } with open('./data/stage2_summary.json', 'w') as f: json.dump(summary, f, indent=2) def main(): """Main execution for Stage 2""" logger.info("Starting Stage 2: Gemini Vision Classification...") # Check API key if not check_api_key(): return # Classify images classify_all_images(max_workers=64) # Reduced to avoid rate limits logger.info("Stage 2 completed successfully!") if __name__ == "__main__": main()