diff --git a/main.py b/main.py index aa7d5b5..0317829 100644 --- a/main.py +++ b/main.py @@ -2,11 +2,7 @@ import cv2 import yt_dlp import numpy as np import json -import time -import base64 -_snapshot_cache = {} -CACHE_TTL = 10 def capture_stream_snapshot(youtube_url, output_file="snapshot.jpg", toBuffer=False): # 1. Configure yt-dlp to get the direct stream URL @@ -33,7 +29,7 @@ def capture_stream_snapshot(youtube_url, output_file="snapshot.jpg", toBuffer=Fa if not cap.isOpened(): print("Error: Could not open video stream.") - return None + return None if toBuffer else None # Read a single frame success, frame = cap.read() @@ -43,126 +39,138 @@ def capture_stream_snapshot(youtube_url, output_file="snapshot.jpg", toBuffer=Fa # Encode frame as JPEG to memory buffer ret, buf = cv2.imencode(".jpg", frame) if ret: + print("Snapshot captured to buffer.") cap.release() return buf.tobytes() else: + print("Error: Could not encode frame to buffer.") cap.release() return None else: # 3. Save the frame as an image cv2.imwrite(output_file, frame) - cap.release() - return True + print(f"Snapshot saved to {output_file}") else: print("Error: Could not read frame from stream.") - cap.release() - return None + + # Cleanup + cap.release() # Shsf Handler def main(args): - route = args.get("route", "") - - # Handle /app/{id}.jpg - if route.startswith("app/") and route.endswith(".jpg"): - url_id = route.split("/")[-1].replace(".jpg", "") - now = time.time() - - if url_id in _snapshot_cache: - ts, encoded = _snapshot_cache[url_id] - if now - ts < CACHE_TTL: - return { - "_shsf": "v1", # Changed to v1 for simpler binary response if v2 is picky - "_code": 200, - "_res": encoded, - "_is_base64": True, - "_headers": { "Content-Type": "image/jpeg", "X-Cache": "HIT" } - } - - # Cache miss - yt_url = f"https://www.youtube.com/watch?v={url_id}" - image_bytes = capture_stream_snapshot(yt_url, toBuffer=True) - if image_bytes: - encoded_content = base64.b64encode(image_bytes).decode('utf-8') - _snapshot_cache[url_id] = (now, encoded_content) - return { - "_shsf": "v1", - "_code": 200, - "_res": encoded_content, - "_is_base64": True, - "_headers": { "Content-Type": "image/jpeg", "X-Cache": "MISS" } - } - + route = args.get("route") if route == "snapshot": url = args.get("body", "") try: url = json.loads(url) except json.JSONDecodeError: + print("Error: Invalid JSON input.") return { "_shsf": "v2", "_code": 400, "_res": {"message": "Invalid JSON input."}, - "_headers": { "Content-Type": "application/json" }, + "_headers": { + "Content-Type": "application/json", + }, } url = url.get("url", "") if not url: + print("Error: URL not provided.") return { "_shsf": "v2", "_code": 400, "_res": {"message": "URL not provided."}, - "_headers": { "Content-Type": "application/json" }, + "_headers": { + "Content-Type": "application/json", + }, } try: - image_bytes = capture_stream_snapshot(url, toBuffer=True) - if image_bytes: - encoded_content = base64.b64encode(image_bytes).decode('utf-8') - return { - "_shsf": "v2", - "_code": 200, - "_res": { "message": "done", "buffer": encoded_content }, - "_headers": { "Content-Type": "application/json" }, - } + capture_stream_snapshot(url, output_file="/tmp/snapshot.jpg") + content = None + with open("/tmp/snapshot.jpg", "rb") as f: + content = f.read() except Exception as e: + print(f"Error: {e}") return { "_shsf": "v2", "_code": 500, - "_res": {"message": str(e)}, - "_headers": { "Content-Type": "application/json" }, + "_res": {"message": f"Could not capture snapshot: {str(e)}"}, + "_headers": { + "Content-Type": "application/json", + }, } + if content is None: + print("Error: Could not capture snapshot.") + return { + "_shsf": "v2", + "_code": 500, + "_res": {"message": "Could not capture snapshot."}, + "_headers": { + "Content-Type": "application/json", + }, + } + + # Convert content to base64 for JSON response + import base64 + encoded_content = base64.b64encode(content).decode('utf-8') + return { "_shsf": "v2", - "_code": 500, - "_res": {"message": "Could not capture snapshot."}, - "_headers": { "Content-Type": "application/json" }, + "_code": 200, + "_res": { + "message": "done", + "buffer": encoded_content + }, + "_headers": { + "Content-Type": "application/json", + }, } - elif route == "health": return { "_shsf": "v2", "_code": 200, "_res": {"message": "Service is healthy."}, - "cache_size": len(_snapshot_cache), - "_headers": { "Content-Type": "application/json" }, + "_headers": { + "Content-Type": "application/json", + }, } - elif route == "default": # UI Route try: + content = "" with open("/app/ui.html", "r") as f: content = f.read() - return { - "_shsf": "v2", - "_code": 200, - "_res": content, - "_headers": { "Content-Type": "text/html" }, - } - except Exception as e: + except FileNotFoundError: + print("Error: /app/ui.html not found.") return { "_shsf": "v2", "_code": 404, - "_res": {"message": "UI file not found or error."}, + "_res": {"message": "UI file not found."}, + "_headers": { + "Content-Type": "application/json", + }, } + except Exception as e: + print(f"Error loading UI: {e}") + return { + "_shsf": "v2", + "_code": 500, + "_res": {"message": f"Server error: {str(e)}"}, + "_headers": { + "Content-Type": "application/json", + }, + } + return { + "_shsf": "v2", + "_code": 200, + "_res": content, + "_headers": { + "Content-Type": "text/html", + }, + } if __name__ == "__main__": - print("running stream-shot") + print("omg im on main") + capture_stream_snapshot("https://www.youtube.com/watch?v=n15V_fCsl_c", output_file="snapshot.jpg") \ No newline at end of file