Create app.py

app.py

@@ -0,0 +1,244 @@
+import os
+import io
+import re
+import json
+import tempfile
+import shutil
+import logging
+import base64
+from concurrent.futures import ThreadPoolExecutor
+
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from fastapi.staticfiles import StaticFiles
+from fastapi.responses import FileResponse
+from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image
+import rarfile
+import zipfile
+import google.generativeai as genai
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI()
+
+# 1. Enable CORS (good practice, though less critical if serving static from same origin)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# 2. Utility Functions
+def parse_srt_time_to_ms(time_str):
+    try:
+        if not time_str: return 0
+        time, ms = time_str.replace(',', '.').split('.')
+        hours, minutes, seconds = map(int, time.split(':'))
+        return (hours * 3600000) + (minutes * 60000) + (seconds * 1000) + int(ms)
+    except Exception:
+        return 0
+
+def parse_filename_to_ms(filename):
+    match = re.search(r'(\d{1,2})_(\d{2})_(\d{2})_(\d{3})', filename)
+    if not match:
+        return None
+    h, m, s, ms = map(int, match.groups())
+    return (h * 3600000) + (m * 60000) + (s * 1000) + ms
+
+def parse_srt(content: str):
+    blocks = content.replace('\r\n', '\n').replace('\r', '\n').strip().split('\n\n')
+    parsed = []
+    for block in blocks:
+        lines = block.split('\n')
+        if len(lines) >= 2:
+            time_line = lines[1]
+            if '-->' in time_line:
+                start_str = time_line.split('-->')[0].strip()
+                text = " ".join(lines[2:]) if len(lines) > 2 else "[BLANK]"
+                parsed.append({
+                    "id": lines[0],
+                    "time": time_line,
+                    "startTimeMs": parse_srt_time_to_ms(start_str),
+                    "text": text
+                })
+    return parsed
+
+def compress_image(image_bytes, quality=70, max_width=800):
+    try:
+        img = Image.open(io.BytesIO(image_bytes))
+        if img.mode != 'RGB':
+            img = img.convert('RGB')
+        
+        width, height = img.size
+        if width > max_width:
+            height = int((height * max_width) / width)
+            width = max_width
+            img = img.resize((width, height), Image.Resampling.LANCZOS)
+        
+        buffer = io.BytesIO()
+        img.save(buffer, format="JPEG", quality=quality)
+        return buffer.getvalue()
+    except Exception as e:
+        logger.error(f"Image compression failed: {e}")
+        return None
+
+def process_batch_gemini(api_key, items):
+    try:
+        genai.configure(api_key=api_key)
+        model = genai.GenerativeModel('gemini-2.0-flash') 
+
+        prompt_parts = [
+            "You are a Subtitle Quality Control (QC) bot.",
+            f"I will provide {len(items)} images and the EXPECTED subtitle text for each.",
+            "Return a JSON array strictly following this schema:",
+            '[{"index": <int>, "detected_text": "<string>", "match": <bool>, "reason": "<string>"}, ...]',
+            "Return ONLY the JSON. No markdown."
+        ]
+
+        for item in items:
+            prompt_parts.append(f"\n--- Item {item['index']} ---")
+            prompt_parts.append(f"Index: {item['index']}")
+            prompt_parts.append(f"Expected Text: \"{item['expected_text']}\"")
+            prompt_parts.append(f"Image:")
+            img = Image.open(io.BytesIO(item['image_data']))
+            prompt_parts.append(img)
+
+        response = model.generate_content(prompt_parts)
+        text = response.text.replace("```json", "").replace("```", "").strip()
+        return json.loads(text)
+    except Exception as e:
+        logger.error(f"Gemini API Error with key ...{api_key[-4:]}: {e}")
+        return None 
+
+# 3. API Endpoint
+@app.post("/api/analyze")
+async def analyze_subtitles(
+    srt_file: UploadFile = File(...),
+    media_files: list[UploadFile] = File(...),
+    api_keys: str = Form(...),
+    batch_size: int = Form(20)
+):
+    temp_dir = tempfile.mkdtemp()
+    try:
+        # Read SRT
+        srt_content = (await srt_file.read()).decode('utf-8', errors='ignore')
+        srt_data = parse_srt(srt_content)
+        srt_data.sort(key=lambda x: x['startTimeMs'])
+
+        # Process Media
+        images = []
+        
+        for file in media_files:
+            file_path = os.path.join(temp_dir, file.filename)
+            with open(file_path, "wb") as f:
+                shutil.copyfileobj(file.file, f)
+            
+            if file.filename.lower().endswith('.rar'):
+                # Docker guarantees 'unrar' is installed
+                try:
+                    with rarfile.RarFile(file_path) as rf:
+                        rf.extractall(temp_dir)
+                except rarfile.RarCannotExec:
+                     raise HTTPException(status_code=500, detail="Unrar executable not found in container.")
+            elif file.filename.lower().endswith('.zip'):
+                with zipfile.ZipFile(file_path, 'r') as zf:
+                    zf.extractall(temp_dir)
+        
+        for root, _, files in os.walk(temp_dir):
+            for filename in files:
+                if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.webp', '.bmp')):
+                    full_path = os.path.join(root, filename)
+                    ms = parse_filename_to_ms(filename)
+                    if ms is not None:
+                        with open(full_path, "rb") as f:
+                            raw_bytes = f.read()
+                            compressed = compress_image(raw_bytes)
+                            if compressed:
+                                images.append({
+                                    "filename": filename,
+                                    "timeMs": ms,
+                                    "data": compressed
+                                })
+
+        images.sort(key=lambda x: x['timeMs'])
+        
+        pairs = []
+        for i in range(len(images)):
+            img = images[i]
+            srt = srt_data[i] if i < len(srt_data) else None
+            
+            if srt:
+                thumb_bytes = compress_image(img['data'], quality=50, max_width=300)
+                thumb_b64 = base64.b64encode(thumb_bytes).decode('utf-8')
+                
+                pairs.append({
+                    "index": i,
+                    "image_data": img['data'],
+                    "expected_text": srt['text'],
+                    "srt_id": srt['id'],
+                    "srt_time": srt['time'],
+                    "filename": img['filename'],
+                    "thumb": f"data:image/jpeg;base64,{thumb_b64}",
+                    "status": "pending"
+                })
+
+        if not pairs:
+            return {"status": "error", "message": "No valid image/subtitle pairs found."}
+
+        # Process Gemini
+        keys = [k.strip() for k in api_keys.split('\n') if k.strip()]
+        if not keys:
+            raise HTTPException(status_code=400, detail="No API Keys provided")
+
+        results_map = {}
+        batches = [pairs[i:i + batch_size] for i in range(0, len(pairs), batch_size)]
+        
+        def worker(batch_idx, batch):
+            key = keys[batch_idx % len(keys)]
+            return process_batch_gemini(key, batch)
+
+        with ThreadPoolExecutor(max_workers=len(keys)) as executor:
+            futures = [executor.submit(worker, i, b) for i, b in enumerate(batches)]
+            for future in futures:
+                res = future.result()
+                if res:
+                    for item in res:
+                        results_map[item['index']] = item
+
+        final_output = []
+        for p in pairs:
+            analysis = results_map.get(p['index'])
+            status = "pending"
+            reason = ""
+            detected = ""
+            if analysis:
+                status = "match" if analysis['match'] else "mismatch"
+                reason = analysis.get('reason', '')
+                detected = analysis.get('detected_text', '')
+
+            final_output.append({
+                "id": p['index'],
+                "filename": p['filename'],
+                "thumb": p['thumb'],
+                "expected": p['expected_text'],
+                "detected": detected,
+                "status": status,
+                "reason": reason,
+                "srt_id": p['srt_id'],
+                "srt_time": p['srt_time']
+            })
+
+        return {"status": "success", "results": final_output}
+
+    except Exception as e:
+        logger.error(f"Server Error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        shutil.rmtree(temp_dir)
+
+# 4. Serve Static Files (Frontend)
+app.mount("/", StaticFiles(directory="static", html=True), name="static")