Upload 12 files

app.py

@@ -1,41 +1,38 @@
-from dotenv import load_dotenv
+import os
 import base64
 import json
-import os
+from dotenv import load_dotenv
+
 load_dotenv(override=True)
-encoded_env = os.getenv("ENCODED_ENV")
-print(f"Encoded environment: {encoded_env}")
+encoded_env = os.getenv("ENCODED_ENV_IMAGE")
 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 torch
-import faulthandler
-import json
 from fastapi import FastAPI
 from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+from dotenv import load_dotenv
+import faulthandler
 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
+from src.utils.zip_utils import extract_zip_file
+from src.utils.model_utils import init_models, search_similar_images
+from src.firebase.firebase_provider import process_images
 
 # 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)
 
-# Load environment variables from base64 encoded string
+# Load environment variables
 
 
 # Initialize FastAPI app
@@ -50,30 +47,21 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Initialize paths and extract data
+# Initialize paths and models
 index_path = "./model/db_vit_b_16.index"
 onnx_path = "./model/vit_b_16_feature_extractor.onnx"
+index, feature_extractor = init_models(index_path, onnx_path)
+
+# Extract images if needed
 zip_file = "./images_2.zip"
 extract_path = "./data"
-
-# Check if index file exists
-if not os.path.exists(index_path):
-    raise FileNotFoundError(f"Index file not found: {index_path}")
-
-try:
-    # Load FAISS index
-    index = faiss.read_index(index_path)
-    print(f"Successfully loaded FAISS index from {index_path}")
-    # Initialize feature extractor with ONNX support
-    feature_extractor = FeatureExtractor(base_model="vit_b_16", onnx_path=onnx_path)
-    print("Successfully initialized feature extractor with ONNX support")
-except Exception as e:
-    raise RuntimeError(f"Error initializing models: {str(e)}")
-
-# Extract zip file if needed
 extract_zip_file(zip_file, extract_path)
 
 
+class ImageSearchBody(BaseModel):
+    base64_image: str = Field(..., title="Base64 Image String")
+
+
 @app.post("/search-image/")
 def search_image(body: ImageSearchBody):
     try:
@@ -81,27 +69,31 @@ def search_image(body: ImageSearchBody):
         image = base64_to_image(body.base64_image)
 
         # Extract features using ONNX model
-        output = feature_extractor.extract_features(image)
-
-        # Prepare features for FAISS search
-        output = output.view(output.size(0), -1)
-        output = output / output.norm(p=2, dim=1, keepdim=True)
+        features = feature_extractor.extract_features(image)
 
         # Search for similar images
-        D, I = index.search(output.cpu().numpy(), 1)
+        D, I = search_similar_images(index, features)
 
         # Get the matched image
-        image_list = sorted([f for f in os.listdir(extract_path + "/images") if is_image_file(f)])
+        image_list = sorted(
+            [f for f in os.listdir(extract_path + "/images") if is_image_file(f)]
+        )
         image_name = image_list[int(I[0][0])]
         matched_image_path = f"{extract_path}/images/{image_name}"
         matched_image = Image.open(matched_image_path)
         matched_image_base64 = image_to_base64(matched_image)
-        
-        # 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
-        
+
+        # Post-process image name: remove underscores, numbers, and file extension
+        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
+
         return JSONResponse(
             content={
                 "image_base64": matched_image_base64,
@@ -114,26 +106,35 @@ 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
         )
 
 
+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: ImageUploadBody):
+async def upload_image(body: Body):
     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)
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)

src/firebase/firebase_provider.py

@@ -8,6 +8,7 @@ import asyncio
 from typing import List, Optional
 from datetime import datetime
 import pytz
+from src.utils.image_utils import base64_to_image
 
 
 import asyncio
@@ -36,20 +37,23 @@ async def upload_file_to_storage(file_path: str, file_name: str) -> str:
     """
     Asynchronous wrapper to upload a file to Firebase Storage using a thread pool.
 
-    param:
+    Args:
         file_path: str - The path of the file on the local machine to be uploaded.
         file_name: str - The name of the file in Firebase Storage.
 
-    return:
+    Returns:
         str - The public URL of the uploaded file.
     """
     loop = asyncio.get_event_loop()
 
-    # Run the synchronous `upload_file_to_storage_sync` in a thread pool.
-    public_url = await loop.run_in_executor(
-        None, functools.partial(upload_file_to_storage_sync, file_path, file_name)
-    )
+    # Run the synchronous upload in a thread pool
+    def upload_sync():
+        blob = firebase_bucket.blob(file_name)
+        blob.upload_from_filename(file_path)
+        blob.make_public()
+        return blob.public_url
 
+    public_url = await loop.run_in_executor(None, upload_sync)
     return public_url
 
 
@@ -81,8 +85,8 @@ def delete_file_by_url(public_url):
     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]
-        
+        file_name = public_url.split("/")[-1]
+
         # Delete the file using the extracted name
         return delete_file_from_storage(file_name)
     except Exception as e:
@@ -121,7 +125,7 @@ def download_file_from_storage(file_name, destination_path):
 
 
 async def upload_base64_image_to_storage(
-    base64_image: str, file_name: str
+    base64_image: str, file_name: str, format: str = "JPEG"
 ) -> Optional[str]:
     """
     Upload a base64 image to Firebase Storage asynchronously.
@@ -129,46 +133,42 @@ async def upload_base64_image_to_storage(
     Args:
         base64_image: str - The base64 encoded image
         file_name: str - The name of the file to be uploaded
+        format: str - The format to save the image in (JPEG, PNG, etc.)
 
     Returns:
         Optional[str] - The public URL of the uploaded file or None if failed
     """
     try:
-        # Run CPU-intensive operations in thread pool
-        loop = asyncio.get_event_loop()
-
-        # Decode base64 in thread pool
-        image_data = await loop.run_in_executor(
-            None, lambda: base64.b64decode(base64_image)
-        )
-
-        # Open and process image in thread pool
-        image = await loop.run_in_executor(
-            None, lambda: Image.open(io.BytesIO(image_data))
-        )
+        # Convert base64 to PIL Image
+        image = base64_to_image(base64_image)
 
-        # Create unique temp file path
+        # Create unique temp file path with appropriate extension
         temp_file_path = os.path.join(
-            tempfile.gettempdir(), f"{file_name}_{datetime.now().timestamp()}.jpg"
+            tempfile.gettempdir(),
+            f"{file_name}_{datetime.now().timestamp()}.{format.lower()}",
         )
 
-        # Save image in thread pool
-        await loop.run_in_executor(
-            None, lambda: image.save(temp_file_path, format="JPEG")
-        )
+        # Save image in the specified format
+        image.save(temp_file_path, format=format)
 
         try:
             # Upload to Firebase
             public_url = await upload_file_to_storage(
-                temp_file_path, f"{file_name}.jpg"
+                temp_file_path, f"{file_name}.{format.lower()}"
             )
             return public_url
         finally:
-            # Clean up temp file in thread pool
-            await loop.run_in_executor(None, os.remove, temp_file_path)
+            # Clean up temp file
+            if os.path.exists(temp_file_path):
+                os.remove(temp_file_path)
 
     except Exception as e:
         print(f"Error processing image {file_name}: {str(e)}")
+        # If format is not JPEG, try again with JPEG
+        if format.upper() != "JPEG":
+            return await upload_base64_image_to_storage(
+                base64_image, file_name, format="JPEG"
+            )
         return None
 
 
@@ -190,6 +190,18 @@ async def process_images(base64_images: List[str]) -> List[Optional[str]]:
             .strftime("%Y-%m-%d_%H-%M-%S")
         )
         file_name = f"image_{timestamp}_{idx}"
-        tasks.append(upload_base64_image_to_storage(base64_image, file_name))
+
+        # Determine format from base64 header or default to JPEG
+        format = "JPEG"
+        if "data:image/" in base64_image:
+            mime_type = base64_image.split(";")[0].split("/")[1]
+            if mime_type == "png":
+                format = "PNG"
+            elif mime_type == "webp":
+                format = "WEBP"
+
+        tasks.append(
+            upload_base64_image_to_storage(base64_image, file_name, format=format)
+        )
 
     return await asyncio.gather(*tasks, return_exceptions=True)

src/utils/image_utils.py

@@ -5,55 +5,105 @@ from fastapi import HTTPException
 
 
 def base64_to_image(base64_str: str) -> Image.Image:
-    """
-    Convert a base64 string to a PIL Image.
+    """Convert base64 string to PIL Image.
     
     Args:
-        base64_str (str): The base64 encoded image string
+        base64_str: Base64 encoded image string
         
     Returns:
-        Image.Image: The decoded PIL Image
+        PIL.Image: Decoded image
         
     Raises:
-        HTTPException: If the base64 string is invalid
+        HTTPException: If base64 string is invalid
     """
     try:
-        # Handle frontend format: data:image/jpeg;base64,{base64_data}
+        # Handle frontend base64 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")
+        image = Image.open(BytesIO(image_data))
+        
+        # Convert RGBA to RGB if necessary
+        if image.mode in ('RGBA', 'LA'):
+            background = Image.new('RGB', image.size, (255, 255, 255))
+            if image.mode == 'RGBA':
+                background.paste(image, mask=image.split()[3])  # 3 is the alpha channel
+            else:
+                background.paste(image, mask=image.split()[1])  # 1 is the alpha channel
+            image = background
+        elif image.mode != 'RGB':
+            image = image.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.
+def image_to_base64(image: Image.Image, format: str = "JPEG") -> str:
+    """Convert PIL Image to base64 string.
     
     Args:
-        image (Image.Image): The PIL Image to convert
+        image: PIL Image object
+        format: Output format (JPEG, PNG, etc.)
         
     Returns:
-        str: The base64 encoded image string
+        str: Base64 encoded image string
     """
-    buffered = BytesIO()
-    image.save(buffered, format="JPEG")
-    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+    try:
+        # Convert RGBA to RGB if saving as JPEG
+        if format.upper() == "JPEG" and image.mode in ('RGBA', 'LA'):
+            background = Image.new('RGB', image.size, (255, 255, 255))
+            if image.mode == 'RGBA':
+                background.paste(image, mask=image.split()[3])
+            else:
+                background.paste(image, mask=image.split()[1])
+            image = background
+        elif format.upper() == "JPEG" and image.mode != 'RGB':
+            image = image.convert('RGB')
+            
+        buffered = BytesIO()
+        image.save(buffered, format=format)
+        return base64.b64encode(buffered.getvalue()).decode("utf-8")
+    except Exception as e:
+        print(f"Error converting image to base64: {str(e)}")
+        # Try JPEG as fallback
+        if format.upper() != "JPEG":
+            return image_to_base64(image, format="JPEG")
+        raise
 
 
 def is_image_file(filename: str) -> bool:
-    """
-    Check if a filename has a valid image extension.
+    """Check if a filename has a valid image extension.
     
     Args:
-        filename (str): The filename to check
+        filename: Name of the file to check
         
     Returns:
-        bool: True if the file has a valid image extension
+        bool: True if file has valid image extension
     """
     valid_extensions = (".png", ".jpg", ".jpeg", ".bmp", ".gif", ".tiff", ".webp")
-    return filename.lower().endswith(valid_extensions) 
+    return filename.lower().endswith(valid_extensions)
+
+
+def get_image_format(filename: str) -> str:
+    """Get the format to use for saving an image based on its filename.
+    
+    Args:
+        filename: Name of the file
+        
+    Returns:
+        str: Format to use (JPEG, PNG, etc.)
+    """
+    ext = filename.lower().split('.')[-1]
+    if ext in ('jpg', 'jpeg'):
+        return 'JPEG'
+    elif ext == 'png':
+        return 'PNG'
+    elif ext == 'webp':
+        return 'WEBP'
+    elif ext == 'gif':
+        return 'GIF'
+    else:
+        return 'JPEG'  # Default to JPEG 

src/utils/model_utils.py

@@ -0,0 +1,62 @@
+import os
+import faiss
+import torch
+from src.modules.feature_extractor import FeatureExtractor
+
+
+def init_models(index_path: str, onnx_path: str) -> tuple[faiss.IndexFlatIP, FeatureExtractor]:
+    """Initialize FAISS index and feature extractor.
+    
+    Args:
+        index_path: Path to FAISS index file
+        onnx_path: Path to ONNX model file
+        
+    Returns:
+        tuple: (FAISS index, Feature extractor)
+        
+    Raises:
+        FileNotFoundError: If index file doesn't exist
+        RuntimeError: If model initialization fails
+    """
+    # Check if index file exists
+    if not os.path.exists(index_path):
+        raise FileNotFoundError(f"Index file not found: {index_path}")
+
+    try:
+        # Load FAISS index
+        index = faiss.read_index(index_path)
+        print(f"Successfully loaded FAISS index from {index_path}")
+        
+        # Initialize feature extractor with ONNX support
+        feature_extractor = FeatureExtractor(base_model="vit_b_16", onnx_path=onnx_path)
+        print("Successfully initialized feature extractor with ONNX support")
+        
+        return index, feature_extractor
+        
+    except Exception as e:
+        raise RuntimeError(f"Error initializing models: {str(e)}")
+
+
+def search_similar_images(
+    index: faiss.IndexFlatIP,
+    features: torch.Tensor,
+    k: int = 1
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Search for similar images using FAISS index.
+    
+    Args:
+        index: FAISS index
+        features: Image features to search for
+        k: Number of similar images to return
+        
+    Returns:
+        tuple: (Distances, Indices)
+    """
+    # Prepare features for FAISS search
+    features = features.view(features.size(0), -1)
+    features = features / features.norm(p=2, dim=1, keepdim=True)
+    
+    # Search for similar images
+    D, I = index.search(features.cpu().numpy(), k)
+    
+    return D, I 

src/utils/zip_utils.py

@@ -0,0 +1,43 @@
+import os
+import zipfile
+
+
+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: Path to the zip file
+        destination_folder: Path to the destination folder
+        
+    Raises:
+        FileNotFoundError: If 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}")