Update app.py

app.py

@@ -144,44 +144,52 @@ async def upload_image(file: UploadFile = File(...)):
 
 
 
-T_REPO_ID = "rahul7star/ohamlab"
-T_FOLDER = "demo"
-T_BASE_URL = f"https://huggingface.co/{T_REPO_ID}/resolve/main/"
+import io
+from fastapi import FastAPI, Query
+from huggingface_hub import list_repo_files, hf_hub_download, upload_file
 
+app = FastAPI()
+
+T_REPO_ID = "rahul7star/ohamlab"
 DESCRIPTION_TEXT = (
     "Ra3hul is wearing a black jacket over a striped white t-shirt with blue jeans. "
     "He is standing near a lake with his arms spread wide open, with mountains and cloudy skies in the background."
 )
 
+def is_image_file(filename: str) -> bool:
+    return filename.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
+
 @app.post("/filter-images")
-def filter_and_rename_images():
+def filter_and_rename_images(folder: str = Query("demo", description="Folder path in repo to scan")):
     try:
         all_files = list_repo_files(T_REPO_ID)
-        folder_prefix = T_FOLDER.rstrip("/") + "/"
-        filter_folder = f"filter-{T_FOLDER}"
+        folder_prefix = folder.rstrip("/") + "/"
+        filter_folder = f"filter-{folder.rstrip('/')}"
         filter_prefix = filter_folder + "/"
 
+        # Filter images only directly in the folder (no subfolders)
         image_files = [
             f for f in all_files
             if f.startswith(folder_prefix)
-            and "/" not in f[len(folder_prefix):]
-            and f.lower().endswith((".png", ".jpg", ".jpeg", ".webp"))
+            and "/" not in f[len(folder_prefix):]  # no deeper path
+            and is_image_file(f)
         ]
 
         if not image_files:
-            return {"error": f"No images found in folder '{T_FOLDER}'"}
+            return {"error": f"No images found in folder '{folder}'"}
 
         uploaded_files = []
 
         for idx, orig_path in enumerate(image_files, start=1):
-            # Download file content into bytes (no local disk)
+            # Download image content bytes (uses local cache)
             local_path = hf_hub_download(repo_id=T_REPO_ID, filename=orig_path)
             with open(local_path, "rb") as f:
                 file_bytes = f.read()
 
-            new_image_name = f"image{idx}.jpeg"  # rename all to .jpeg
+            # Rename images as image1.jpeg, image2.jpeg, ...
+            new_image_name = f"image{idx}.jpeg"
 
-            # Upload image bytes directly from memory
+            # Upload renamed image from memory
             upload_file(
                 path_or_fileobj=io.BytesIO(file_bytes),
                 path_in_repo=filter_prefix + new_image_name,
@@ -192,7 +200,7 @@ def filter_and_rename_images():
             )
             uploaded_files.append(filter_prefix + new_image_name)
 
-            # Create and upload text file from string bytes
+            # Create and upload text file for each image
             txt_filename = f"image{idx}.txt"
             upload_file(
                 path_or_fileobj=io.BytesIO(DESCRIPTION_TEXT.encode("utf-8")),
@@ -204,7 +212,10 @@ def filter_and_rename_images():
             )
             uploaded_files.append(filter_prefix + txt_filename)
 
-        return {"message": f"Processed and uploaded {len(image_files)} images and text files.", "files": uploaded_files}
+        return {
+            "message": f"Processed and uploaded {len(image_files)} images and text files.",
+            "files": uploaded_files,
+        }
 
     except Exception as e:
         return {"error": str(e)}