Update main.py

main.py

@@ -8,6 +8,12 @@ import cv2
 import numpy as np
 import base64
 import requests
+
+import mimetypes
+import tempfile
+import subprocess
+
+
 from PIL import Image
 from io import BytesIO
 from tensorflow.keras.applications import MobileNetV2
@@ -70,17 +76,58 @@ def mobilenet_similarity(img1, img2):
 def load_image(source):
     Image.MAX_IMAGE_PIXELS = None
 
+    def extract_frame_from_video(video_path_or_url, time_sec):
+        with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as temp_frame:
+            frame_path = temp_frame.name
+
+        command = [
+            ffmpeg_path,
+            "-ss", str(time_sec),
+            "-i", video_path_or_url,
+            "-frames:v", "1",
+            "-q:v", "2",
+            "-y",
+            frame_path
+        ]
+        subprocess.run(command, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+
+        if not os.path.exists(frame_path):
+            raise ValueError("Failed to extract frame from video.")
+
+        frame = cv2.imread(frame_path, cv2.IMREAD_GRAYSCALE)
+        os.remove(frame_path)
+        return frame
+
     if source.startswith('http'):
+        mime_type, _ = mimetypes.guess_type(source)
+        if mime_type and mime_type.startswith('video'):
+            return extract_frame_from_video(source, frame_time)
+
+        # Assume imagem
         response = requests.get(source)
         img = np.asarray(bytearray(response.content), dtype=np.uint8)
         img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)
-    else:
-        img = base64.b64decode(source)
-        img = Image.open(BytesIO(img))
-        img = np.array(img)
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+        return img
 
-    return img
+    else:
+        try:
+            img_bytes = base64.b64decode(source)
+            mime_type = mimetypes.guess_type("data")[0]  # fallback
+
+            # Tenta abrir como imagem
+            img = Image.open(BytesIO(img_bytes))
+            img = np.array(img)
+            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+            return img
+        except Exception:
+            # Se falhar, assumimos que é vídeo em base64 e extraímos frame
+            with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_video:
+                temp_video.write(base64.b64decode(source))
+                temp_video_path = temp_video.name
+
+            frame = extract_frame_from_video(temp_video_path, frame_time)
+            os.remove(temp_video_path)
+            return frame
 
 app = FastAPI()