Update similarity.py (#6)

- Update similarity.py (24e1c6ada9903aaf125e34b5b82c83b8ae3c0e7b)

similarity.py

@@ -73,20 +73,90 @@ def ssim_sim(img1, img2):
     return (s + 1) * 50
 
 
-def load_image_url(source):
+def load_image(source, assetCode, contentType=None, ffmpeg_path='ffmpeg', frame_time=1):
     Image.MAX_IMAGE_PIXELS = None
 
-    if source.startswith('http'):
-        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)
+    def extract_frame_from_video(video_path_or_url, time_sec):
+        print(f"[INFO] A extrair frame do vídeo: {video_path_or_url} no segundo {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
+        ]
+
+        print(f"[DEBUG] Comando ffmpeg: {' '.join(command)}")
+
+        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+        if result.returncode != 0:
+            print(f"[ERRO] ffmpeg falhou com código {result.returncode}")
+            print(f"[ERRO] stderr: {result.stderr.decode('utf-8')}")
+            raise RuntimeError("Erro ao extrair frame com ffmpeg.")
+
+        if not os.path.exists(frame_path):
+            print("[ERRO] Frame não criado. Verifica se o caminho do vídeo está correto e acessível.")
+            raise ValueError("Frame não encontrado após execução do ffmpeg.")
 
-    return img
+        frame = cv2.imread(frame_path, cv2.IMREAD_GRAYSCALE)
+        os.remove(frame_path)
+
+        if frame is None:
+            print("[ERRO] Falha ao ler frame extraído com OpenCV.")
+            raise ValueError("Erro ao carregar frame extraído.")
+
+        print(f"[SUCESSO] Frame extraído com sucesso de {video_path_or_url}")
+        return frame
+
+    try:
+        if source.startswith('http'):
+            print(f"[INFO] Content-Type de {assetCode} é {contentType}")
+
+            if contentType and contentType.startswith('video'):
+                return extract_frame_from_video(source, frame_time)
+
+            print(f"[INFO] A carregar imagem {assetCode} a partir de URL")
+            response = requests.get(source)
+            img = np.asarray(bytearray(response.content), dtype=np.uint8)
+            img = cv2.imdecode(img, cv2.IMREAD_GRAYSCALE)
+            return img
+
+        else:
+            print(f"[INFO] A tentar carregar base64 de {assetCode} como imagem ou vídeo.")
+
+            try:
+                img_bytes = base64.b64decode(source)
+
+                if contentType and contentType.startswith('image'):
+                    print(f"[INFO] Base64 de {assetCode} identificado como imagem")
+                    img = Image.open(BytesIO(img_bytes))
+                    img = np.array(img)
+                    img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+                    return img
+                else:
+                    print(f"[INFO] Base64 de {assetCode} identificado como vídeo")
+                    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_video:
+                        temp_video.write(img_bytes)
+                        temp_video_path = temp_video.name
+
+                    frame = extract_frame_from_video(temp_video_path, frame_time)
+                    os.remove(temp_video_path)
+                    return frame
+
+            except Exception as e:
+                print(f"[ERRO] Falha ao processar base64 de {assetCode}: {e}")
+                raise
+
+    except Exception as e:
+        print(f"[ERRO] Falha ao carregar imagem para {assetCode}: {e}")
+        return None
 
 def check_similarity(images: List[RequestModel]):
     logging.info(f"Checking similarity for main source with resource id {images[0].originId}")