Update app.py

app.py

@@ -1,149 +1,33 @@
-from fastapi import FastAPI, HTTPException, BackgroundTasks
-from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
-import os
+from fastapi import FastAPI, UploadFile, File
+from fastapi.responses import JSONResponse
+from fastapi.staticfiles import StaticFiles
+from ultralytics import YOLO
+import shutil
 import uuid
-import logging
+import base64
+import cv2
 
-from pytube import YouTube
-from moviepy.editor import AudioFileClip
-import whisper
-from googletrans import Translator
-from gtts import gTTS
+app = FastAPI()
 
-# Setup
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+# Mount static frontend
+app.mount("/", StaticFiles(directory="static", html=True), name="static")
 
-app = FastAPI(title="YouTube Streaming Translator API (No Cookies, No ffmpeg)")
+# Load YOLOv8 model
+model = YOLO("yolov8n.pt")  # Use your own trained model if needed
 
-YOUTUBE_DIR = "/tmp/youtube_translator"
-os.makedirs(YOUTUBE_DIR, exist_ok=True)
+@app.post("/detect")
+async def detect(file: UploadFile = File(...)):
+    # Save image to disk
+    file_name = f"uploads/{uuid.uuid4()}.jpg"
+    with open(file_name, "wb") as buffer:
+        shutil.copyfileobj(file.file, buffer)
 
-# Load whisper model (use "tiny" for speed)
-try:
-    model = whisper.load_model("tiny")
-    logger.info("Whisper model loaded successfully")
-except Exception as e:
-    logger.error(f"Failed to load whisper model: {e}")
-    model = None
+    # Detect objects
+    results = model(file_name)
+    result_image = results[0].plot()  # Get image with boxes
 
-translator = Translator()
+    # Encode image to base64 to send back
+    _, buffer = cv2.imencode('.jpg', result_image)
+    b64_encoded = base64.b64encode(buffer).decode('utf-8')
 
-class VideoRequest(BaseModel):
-    url: str
-    timestamp: int = 0
-    chunk_size: int = 15
-    target_language: str = "en"
-
-@app.post("/process-chunk/")
-async def process_chunk(request: VideoRequest, background_tasks: BackgroundTasks):
-    request_id = str(uuid.uuid4())
-    audio_path = os.path.join(YOUTUBE_DIR, f"{request_id}.mp4")
-    chunk_path = os.path.join(YOUTUBE_DIR, f"{request_id}_chunk.mp3")
-
-    try:
-        # Download audio using pytube
-        yt = YouTube(request.url)
-        stream = yt.streams.filter(only_audio=True).first()
-        stream.download(output_path=YOUTUBE_DIR, filename=f"{request_id}.mp4")
-
-        # Extract audio chunk using moviepy
-        with AudioFileClip(audio_path) as audio:
-            start = request.timestamp
-            end = start + request.chunk_size
-            audio.subclip(start, end).write_audiofile(chunk_path, codec='mp3')
-
-        # Process audio chunk in background
-        background_tasks.add_task(
-            process_audio_chunk,
-            chunk_path,
-            request.target_language,
-            request_id
-        )
-
-        return {"request_id": request_id, "status": "processing"}
-    except Exception as e:
-        logger.error(f"Error processing chunk: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-    finally:
-        if os.path.exists(audio_path):
-            os.remove(audio_path)
-
-async def process_audio_chunk(chunk_path, target_language, request_id):
-    try:
-        # Step 1: Transcribe
-        logger.info(f"Transcribing audio chunk: {chunk_path}")
-        result = model.transcribe(chunk_path)
-        transcription = result["text"]
-
-        # Step 2: Translate
-        logger.info(f"Translating text to {target_language}")
-        translation = translator.translate(transcription, dest=target_language).text
-
-        # Step 3: TTS
-        logger.info(f"Converting translation to speech")
-        tts_output_path = os.path.join(YOUTUBE_DIR, f"{request_id}_tts.mp3")
-        tts = gTTS(text=translation, lang=target_language)
-        tts.save(tts_output_path)
-
-        # Save translation text
-        text_output_path = os.path.join(YOUTUBE_DIR, f"{request_id}_text.txt")
-        with open(text_output_path, "w", encoding="utf-8") as f:
-            f.write(translation)
-
-        logger.info(f"Audio processing completed for request {request_id}")
-    except Exception as e:
-        logger.error(f"Error processing audio chunk: {e}")
-    finally:
-        if os.path.exists(chunk_path):
-            os.remove(chunk_path)
-
-@app.get("/get-audio/{request_id}")
-async def get_audio(request_id: str):
-    tts_output_path = os.path.join(YOUTUBE_DIR, f"{request_id}_tts.mp3")
-    if not os.path.exists(tts_output_path):
-        raise HTTPException(status_code=404, detail="Audio processing not completed yet or request ID invalid")
-
-    def iterfile():
-        with open(tts_output_path, "rb") as f:
-            yield from f
-
-    return StreamingResponse(
-        iterfile(),
-        media_type="audio/mpeg",
-        headers={"Content-Disposition": f"attachment; filename={request_id}.mp3"}
-    )
-
-@app.get("/get-translation/{request_id}")
-async def get_translation(request_id: str):
-    text_output_path = os.path.join(YOUTUBE_DIR, f"{request_id}_text.txt")
-    if not os.path.exists(text_output_path):
-        raise HTTPException(status_code=404, detail="Translation text not found or processing not completed")
-    with open(text_output_path, "r", encoding="utf-8") as f:
-        translation = f.read()
-    return {"request_id": request_id, "translation": translation}
-
-@app.get("/status/{request_id}")
-async def check_status(request_id: str):
-    tts_output_path = os.path.join(YOUTUBE_DIR, f"{request_id}_tts.mp3")
-    if os.path.exists(tts_output_path):
-        return {"status": "completed", "request_id": request_id}
-    else:
-        chunk_path = os.path.join(YOUTUBE_DIR, f"{request_id}_chunk.mp3")
-        if os.path.exists(chunk_path):
-            return {"status": "processing", "request_id": request_id}
-        else:
-            raise HTTPException(status_code=404, detail="Request ID not found")
-
-@app.get("/")
-async def root():
-    return {"message": "YouTube Streaming Translator API (No Cookies, No ffmpeg)"}
-
-@app.get("/health")
-async def health_check():
-    return {"status": "healthy"}
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+    return JSONResponse(content={"image": b64_encoded})