Update main.py

main.py

@@ -11,7 +11,7 @@ import requests
 
 import mimetypes
 import tempfile
-from moviepy.editor import VideoFileClip
+import subprocess
 
 
 from PIL import Image
@@ -73,17 +73,30 @@ def mobilenet_similarity(img1, img2):
         print("Erro ao calcular similaridade com MobileNet")
         return 0
 
-def load_image(source, frame_time=1):
+def load_image(source, ffmpeg_path='ffmpeg', frame_time=1):
     Image.MAX_IMAGE_PIXELS = None
 
     def extract_frame_from_video(video_path_or_url, time_sec):
-        clip = VideoFileClip(video_path_or_url)
-        frame = clip.get_frame(time_sec)  # frame RGB numpy array
-        clip.reader.close()
-        if clip.audio:
-            clip.audio.reader.close_proc()
-        frame_gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
-        return frame_gray
+        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)
@@ -99,13 +112,15 @@ def load_image(source, frame_time=1):
     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, assume que é vídeo em base64 e extrai frame
+            # 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