Update app.py

app.py

@@ -566,33 +566,51 @@ async def extract_and_store_frames(file_name: str = Form(...)):
 @app.get("/get-video-frames/")
 async def get_video_frames(file_name: str, frame_id: int) -> dict:
     try:
-        # URL du fichier image dans le dataset
-        image_url = f"https://huggingface.co/datasets/{DATASET_REPO}/tree/main/frames/{Path(file_name).stem}/{Path(file_name).stem}_{frame_id}.png"
-        print(image_url)
+        frames = []
+        csv_filename = Path(file_name).stem + ".csv"
+        csv_url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/csv/{csv_filename}"
+        print(csv_url)
         headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-        print(HF_TOKEN)
-        response = requests.head(image_url, headers=headers)
-
-        if response.status_code != 200:
-            raise HTTPException(status_code=404, detail="Image introuvable dans le dataset")
         
-        csv_filename = Path(file_name).stem + ".csv"
-        csv_url = f"https://huggingface.co/datasets/{DATASET_REPO}/tree/main/csv/{csv_filename}"
+        # Récupérer la liste des fichiers dans le dossier frames/{file_name}/
+        dataset_url = f"https://huggingface.co/api/datasets/{DATASET_REPO}/resolve/main/frames/{Path(file_name).stem}/"
+        print(dataset_url)
+        file_list_response = requests.get(dataset_url, headers=headers)
+        available_frames = set()
+        
+        if file_list_response.status_code == 200:
+            file_list = file_list_response.json()
+            available_frames = {file["path"] for file in file_list}
+        
+        # Charger le CSV des annotations si disponible
         csv_response = requests.get(csv_url, headers=headers)
         df_annotations = None
-        
+
         if csv_response.status_code == 200:
             csv_bytes = io.BytesIO(csv_response.content)
             df_annotations = pd.read_csv(csv_bytes)
+
+        for i in range(frame_id, frame_id + 30):
+            frame_filename = f"frames/{Path(file_name).stem}/{Path(file_name).stem}_{i}.png"
+            if frame_filename not in available_frames:
+                continue
+
+            image_url = f"https://huggingface.co/datasets/{DATASET_REPO}/resolve/main/{frame_filename}"
+            class_label, player = None, None
+            
+            if df_annotations is not None and i in df_annotations["frame"].values:
+                row = df_annotations[df_annotations["frame"] == i].iloc[0]
+                class_label, player = row.get("class", None), row.get("player", None)
+
+            frames.append({"frame_index": i, "image_url": image_url, "class": class_label, "player": player})
+
+        if not frames:
+            raise HTTPException(status_code=404, detail="Aucune image trouvée dans la plage demandée")
         
-        class_label, player = None, None
-        if df_annotations is not None and frame_id in df_annotations["frame"].values:
-            row = df_annotations[df_annotations["frame"] == frame_id].iloc[0]
-            class_label, player = row.get("class", None), row.get("player", None)
-        
-        return {"status": "success", "frame_index": frame_id, "image_url": image_url, "class": class_label, "player": player}
+        return {"status": "success", "frames": frames}
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Erreur lors de la récupération de l'image : {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Erreur lors de la récupération des images : {str(e)}")
+