File size: 11,009 Bytes
b72fefd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
#!/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()