Upload 11 files

app.py

@@ -1,39 +1,44 @@
+from dotenv import load_dotenv
+import base64
+import json
+import os
+load_dotenv(override=True)
+encoded_env = os.getenv("ENCODED_ENV")
+print(f"Encoded environment: {encoded_env}")
+if encoded_env:
+    decoded_env = base64.b64decode(encoded_env).decode()
+    env_data = json.loads(decoded_env)
+    for key, value in env_data.items():
+        os.environ[key] = value
+        print(f"Environment variable {key} set to {value}")
+
 import os
 import faiss
-import base64
-from PIL import Image
-from fastapi import FastAPI, HTTPException
+import torch
+import faulthandler
+import json
+from fastapi import FastAPI
 from fastapi.responses import JSONResponse
-from io import BytesIO
-from src.modules.feature_extractor import FeatureExtractor
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, Field
-import json
-from dotenv import load_dotenv
-import faulthandler
-import torch
+from PIL import Image
+
+from src.modules.feature_extractor import FeatureExtractor
+from src.firebase.firebase_provider import process_images
+from src.utils.image_utils import base64_to_image, image_to_base64, is_image_file
+from src.utils.file_utils import extract_zip_file
+from src.models.schemas import ImageSearchBody, ImageUploadBody
 
 # Enable fault handler to debug segmentation faults
 faulthandler.enable()
-load_dotenv(override=True)
 
 # Force CPU mode to avoid segmentation faults with ONNX/PyTorch
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 torch.set_num_threads(1)
 
-encoded_env = os.getenv("ENCODED_ENV")
-if encoded_env:
-    # Decode the base64 string
-    decoded_env = base64.b64decode(encoded_env).decode()
-
-    # Load it as a dictionary
-    env_data = json.loads(decoded_env)
-
-    # Set environment variables
-    for key, value in env_data.items():
-        os.environ[key] = value
+# Load environment variables from base64 encoded string
 
 
+# Initialize FastAPI app
 app = FastAPI(docs_url="/")
 origins = ["*"]
 
@@ -45,9 +50,11 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Initialize paths
+# Initialize paths and extract data
 index_path = "./model/db_vit_b_16.index"
 onnx_path = "./model/vit_b_16_feature_extractor.onnx"
+zip_file = "./images_2.zip"
+extract_path = "./data"
 
 # Check if index file exists
 if not os.path.exists(index_path):
@@ -63,82 +70,10 @@ try:
 except Exception as e:
     raise RuntimeError(f"Error initializing models: {str(e)}")
 
-
-def base64_to_image(base64_str: str) -> Image.Image:
-    try:
-        image_data = base64.b64decode(base64_str)
-        image = Image.open(BytesIO(image_data)).convert("RGB")
-        return image
-    except Exception as e:
-        raise HTTPException(status_code=400, detail="Invalid Base64 image")
-
-
-def image_to_base64(image: Image.Image) -> str:
-    buffered = BytesIO()
-    image.save(buffered, format="JPEG")
-    return base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-
-def extract_zip_file(zip_file_path, destination_folder):
-    """
-    Extract a zip file to a destination folder
-    If destination folder already exists, extraction is skipped
-    
-    Args:
-        zip_file_path: str, path to the zip file
-        destination_folder: str, path to the destination folder
-    
-    Returns:
-        None
-    """
-    import zipfile
-    import os
-    
-    # Check if destination folder already exists
-    if os.path.exists(destination_folder):
-        print(f"Destination folder {destination_folder} already exists. Skipping extraction.")
-        return
-    
-    # Check if zip file exists
-    if not os.path.exists(zip_file_path):
-        raise FileNotFoundError(f"Zip file not found: {zip_file_path}")
-    
-    # Create destination folder
-    os.makedirs(destination_folder, exist_ok=True)
-    
-    # Extract the zip file
-    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
-        for member in zip_ref.infolist():
-            # Handle non-ASCII filenames
-            filename = member.filename.encode('cp437').decode('utf-8')
-            extracted_path = os.path.join(destination_folder, filename)
-            
-            # Create directories if needed
-            os.makedirs(os.path.dirname(extracted_path), exist_ok=True)
-            
-            # Extract file
-            if not filename.endswith('/'):  # Skip directories
-                with zip_ref.open(member) as source, open(extracted_path, 'wb') as target:
-                    target.write(source.read())
-    
-    print(f"Successfully extracted {zip_file_path} to {destination_folder}")
-
-
-zip_file = "./images_2.zip"
-extract_path = "./data"
-# Fix function name to match the defined function
+# Extract zip file if needed
 extract_zip_file(zip_file, extract_path)
 
 
-def is_image_file(filename):
-    valid_extensions = (".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tiff", ".webp")
-    return filename.lower().endswith(valid_extensions)
-
-
-class ImageSearchBody(BaseModel):
-    base64_image: str = Field(..., title="Base64 Image String")
-
-
 @app.post("/search-image/")
 def search_image(body: ImageSearchBody):
     try:
@@ -162,7 +97,7 @@ def search_image(body: ImageSearchBody):
         matched_image = Image.open(matched_image_path)
         matched_image_base64 = image_to_base64(matched_image)
         
-        # Post-process image name: remove underscores, numbers, and file extension
+        # Post-process image name
         image_name_post_process = image_name.replace("_", " ")  # Replace underscores with spaces
         image_name_post_process = ''.join([c for c in image_name_post_process if not c.isdigit()])  # Remove numbers
         image_name_post_process = image_name_post_process.rsplit('.', 1)[0]  # Remove file extension
@@ -179,38 +114,26 @@ def search_image(body: ImageSearchBody):
     except Exception as e:
         print(f"Error in search_image: {str(e)}")
         return JSONResponse(
-            content={"error": f"Error processing image: {str(e)}"}, status_code=500
+            content={"error": f"Error processing image: {str(e)}"}, 
+            status_code=500
         )
 
 
-from src.firebase.firebase_provider import process_images
-
-
-class Body(BaseModel):
-    base64_image: list[str] = Field(..., title="Base64 Image String")
-    model_config = {
-        "json_schema_extra": {
-            "examples": [
-                {
-                    "base64_image": [
-                        "iVBORw0KGgoAAAANSUhEUgAAABQAAAAUCAYAAACNiR0NAAABdUlEQVR42mNk",
-                    ]
-                }
-            ]
-        }
-    }
-
-
 @app.post("/upload_image")
-async def upload_image(body: Body):
+async def upload_image(body: ImageUploadBody):
     try:
         public_url = await process_images(body.base64_image)
-        return JSONResponse(content={"public_url": public_url}, status_code=200)
+        return JSONResponse(
+            content={"public_url": public_url}, 
+            status_code=200
+        )
     except Exception as e:
-        return JSONResponse(content={"error": str(e)}, status_code=500)
+        return JSONResponse(
+            content={"error": str(e)}, 
+            status_code=500
+        )
 
 
 if __name__ == "__main__":
     import uvicorn
-
     uvicorn.run(app, host="0.0.0.0", port=8000)

src/firebase/firebase_provider.py

@@ -70,6 +70,26 @@ def delete_file_from_storage(file_name):
         return False
 
 
+def delete_file_by_url(public_url):
+    """
+    Delete a file from Firebase Storage using its public URL
+    param:
+        public_url: str - The public URL of the file to be deleted
+    return:
+        bool - True if the file is deleted successfully, False if the file is not found
+    """
+    try:
+        # Extract the file name from the public URL
+        # URL format is typically: https://storage.googleapis.com/BUCKET_NAME/FILE_NAME
+        file_name = public_url.split('/')[-1]
+        
+        # Delete the file using the extracted name
+        return delete_file_from_storage(file_name)
+    except Exception as e:
+        print(f"Error deleting file by URL: {e}")
+        return False
+
+
 def list_all_files_in_storage():
     """
     View all files in Firebase Storage

src/models/schemas.py

@@ -0,0 +1,36 @@
+from pydantic import BaseModel, Field
+from typing import List
+
+
+class ImageSearchBody(BaseModel):
+    """
+    Schema for image search request body.
+    """
+    base64_image: str = Field(
+        ..., 
+        title="Base64 Image String",
+        description="Base64 encoded image string to search for similar images"
+    )
+
+
+class ImageUploadBody(BaseModel):
+    """
+    Schema for image upload request body.
+    """
+    base64_image: List[str] = Field(
+        ..., 
+        title="Base64 Image String",
+        description="List of base64 encoded image strings to upload"
+    )
+    
+    model_config = {
+        "json_schema_extra": {
+            "examples": [
+                {
+                    "base64_image": [
+                        "iVBORw0KGgoAAAANSUhEUgAAABQAAAAUCAYAAACNiR0NAAABdUlEQVR42mNk",
+                    ]
+                }
+            ]
+        }
+    } 

src/utils/file_utils.py

@@ -0,0 +1,71 @@
+import os
+import zipfile
+from typing import Optional
+
+
+def extract_zip_file(zip_file_path: str, destination_folder: str) -> None:
+    """
+    Extract a zip file to a destination folder.
+    If destination folder already exists, extraction is skipped.
+    
+    Args:
+        zip_file_path (str): Path to the zip file
+        destination_folder (str): Path to the destination folder
+    
+    Raises:
+        FileNotFoundError: If the zip file doesn't exist
+    """
+    # Check if destination folder already exists
+    if os.path.exists(destination_folder):
+        print(f"Destination folder {destination_folder} already exists. Skipping extraction.")
+        return
+    
+    # Check if zip file exists
+    if not os.path.exists(zip_file_path):
+        raise FileNotFoundError(f"Zip file not found: {zip_file_path}")
+    
+    # Create destination folder
+    os.makedirs(destination_folder, exist_ok=True)
+    
+    # Extract the zip file
+    with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+        for member in zip_ref.infolist():
+            # Handle non-ASCII filenames
+            filename = member.filename.encode('cp437').decode('utf-8')
+            extracted_path = os.path.join(destination_folder, filename)
+            
+            # Create directories if needed
+            os.makedirs(os.path.dirname(extracted_path), exist_ok=True)
+            
+            # Extract file
+            if not filename.endswith('/'):  # Skip directories
+                with zip_ref.open(member) as source, open(extracted_path, 'wb') as target:
+                    target.write(source.read())
+    
+    print(f"Successfully extracted {zip_file_path} to {destination_folder}")
+
+
+def ensure_directory_exists(directory_path: str) -> None:
+    """
+    Ensure that a directory exists, create it if it doesn't.
+    
+    Args:
+        directory_path (str): Path to the directory
+    """
+    os.makedirs(directory_path, exist_ok=True)
+
+
+def get_file_extension(filename: str) -> Optional[str]:
+    """
+    Get the extension of a file.
+    
+    Args:
+        filename (str): Name of the file
+    
+    Returns:
+        Optional[str]: The file extension without the dot, or None if no extension
+    """
+    split = os.path.splitext(filename)
+    if len(split) > 1:
+        return split[1][1:]  # Remove the dot
+    return None 

src/utils/image_utils.py

@@ -0,0 +1,59 @@
+import base64
+from PIL import Image
+from io import BytesIO
+from fastapi import HTTPException
+
+
+def base64_to_image(base64_str: str) -> Image.Image:
+    """
+    Convert a base64 string to a PIL Image.
+    
+    Args:
+        base64_str (str): The base64 encoded image string
+        
+    Returns:
+        Image.Image: The decoded PIL Image
+        
+    Raises:
+        HTTPException: If the base64 string is invalid
+    """
+    try:
+        # Handle frontend format: data:image/jpeg;base64,{base64_data}
+        if "," in base64_str:
+            base64_str = base64_str.split(",", 1)[1]
+            
+        image_data = base64.b64decode(base64_str)
+        image = Image.open(BytesIO(image_data)).convert("RGB")
+        return image
+    except Exception as e:
+        print(f"Base64 decoding error: {str(e)}")
+        raise HTTPException(status_code=400, detail=f"Invalid Base64 image: {str(e)}")
+
+
+def image_to_base64(image: Image.Image) -> str:
+    """
+    Convert a PIL Image to a base64 string.
+    
+    Args:
+        image (Image.Image): The PIL Image to convert
+        
+    Returns:
+        str: The base64 encoded image string
+    """
+    buffered = BytesIO()
+    image.save(buffered, format="JPEG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+
+def is_image_file(filename: str) -> bool:
+    """
+    Check if a filename has a valid image extension.
+    
+    Args:
+        filename (str): The filename to check
+        
+    Returns:
+        bool: True if the file has a valid image extension
+    """
+    valid_extensions = (".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tiff", ".webp")
+    return filename.lower().endswith(valid_extensions)