Upload 7 files

src/config/settings.py

@@ -0,0 +1,18 @@
+WORK_DIR = "./"
+DATA_DIR = f"{WORK_DIR}/data"
+IMAGES_DIR = f"{DATA_DIR}/images"
+RESULTS_DIR = f"{WORK_DIR}/results"
+
+# supported feature extractor models
+FEATURE_EXTRACTOR_MODELS = [
+    "resnet18",
+    "resnet34",
+    "resnet50",
+    "resnet101",
+    "resnet152",
+    "vit_b_16",
+    "vit_b_32",
+    "vit_l_16",
+    "vit_l_32",
+    "vit_h_14",
+]

src/firebase/firebase_config.py

@@ -0,0 +1,36 @@
+import firebase_admin
+from firebase_admin import credentials
+from firebase_admin import storage
+
+import os
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+firebase_url_storageBucket = os.getenv("URL_STORAGEBUCKET")
+
+# Get credentials from environment variables
+credential_firebase = {
+    "type": os.getenv("TYPE"),
+    "project_id": os.getenv("PROJECT_ID"),
+    "private_key_id": os.getenv("PRIVATE_KEY_ID"),
+    "private_key": os.getenv("PRIVATE_KEY"),
+    "client_email": os.getenv("CLIENT_EMAIL"),
+    "client_id": os.getenv("CLIENT_ID"),
+    "auth_uri": os.getenv("AUTH_URI"),
+    "token_uri": os.getenv("TOKEN_URI"),
+    "auth_provider_x509_cert_url": os.getenv("AUTH_PROVIDER_X509_CERT_URL"),
+    "client_x509_cert_url": os.getenv("CLIENT_X509_CERT_URL"),
+    "universe_domain": os.getenv("UNIVERSE_DOMAIN"),
+}
+
+
+# Check if the app is not initialized yet
+if not firebase_admin._apps:
+    # Initialize the app with the credentials
+    cred = credentials.Certificate(credential_firebase)
+    firebase_admin.initialize_app(cred, {"storageBucket": firebase_url_storageBucket})
+
+# Initialize Firestore
+firebase_bucket = storage.bucket(app=firebase_admin.get_app())
+print("Storage connected")

src/firebase/firebase_provider.py

@@ -0,0 +1,175 @@
+from .firebase_config import firebase_bucket
+import base64
+import os
+import tempfile
+from PIL import Image
+import io
+import asyncio
+from typing import List, Optional
+from datetime import datetime
+import pytz
+
+
+import asyncio
+import functools
+
+
+def upload_file_to_storage_sync(file_path, file_name):
+    """
+    Synchronous function to upload a file to Firebase Storage.
+
+    param:
+        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:
+        str - The public URL of the uploaded file.
+    """
+    blob = firebase_bucket.blob(file_name)
+    blob.upload_from_filename(file_path)
+    blob.make_public()
+
+    return blob.public_url
+
+
+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:
+        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:
+        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)
+    )
+
+    return public_url
+
+
+def delete_file_from_storage(file_name):
+    """
+    Delete a file from Firebase Storage
+    param:
+        file_name: str - The name of the file to be deleted
+    return:
+        bool - True if the file is deleted successfully, False if the file is not found
+    """
+    try:
+        blob = firebase_bucket.blob(file_name)
+        blob.delete()
+        return True
+    except Exception as e:
+        print("Error:", e)
+        return False
+
+
+def list_all_files_in_storage():
+    """
+    View all files in Firebase Storage
+    return:
+        dict - Dictionary with keys are names and values are url of all files in Firebase Storage
+    """
+    blobs = firebase_bucket.list_blobs()
+    blob_dict = {blob.name: blob.public_url for blob in blobs}
+    return blob_dict
+
+
+def download_file_from_storage(file_name, destination_path):
+    """
+    Download a file from Firebase Storage
+    param:
+        file_name: str - The name of the file to be downloaded
+        destination_path: str - The path to save the downloaded file
+    return:
+        bool - True if the file is downloaded successfully, False if the file is not found
+    """
+    try:
+        blob = firebase_bucket.blob(file_name)
+        blob.download_to_filename(destination_path)
+        print("da tai xun thanh cong")
+        return True
+    except Exception as e:
+        print("Error:", e)
+        return False
+
+
+async def upload_base64_image_to_storage(
+    base64_image: str, file_name: str
+) -> Optional[str]:
+    """
+    Upload a base64 image to Firebase Storage asynchronously.
+
+    Args:
+        base64_image: str - The base64 encoded image
+        file_name: str - The name of the file to be uploaded
+
+    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))
+        )
+
+        # Create unique temp file path
+        temp_file_path = os.path.join(
+            tempfile.gettempdir(), f"{file_name}_{datetime.now().timestamp()}.jpg"
+        )
+
+        # Save image in thread pool
+        await loop.run_in_executor(
+            None, lambda: image.save(temp_file_path, format="JPEG")
+        )
+
+        try:
+            # Upload to Firebase
+            public_url = await upload_file_to_storage(
+                temp_file_path, f"{file_name}.jpg"
+            )
+            return public_url
+        finally:
+            # Clean up temp file in thread pool
+            await loop.run_in_executor(None, os.remove, temp_file_path)
+
+    except Exception as e:
+        print(f"Error processing image {file_name}: {str(e)}")
+        return None
+
+
+async def process_images(base64_images: List[str]) -> List[Optional[str]]:
+    """
+    Process multiple base64 images concurrently.
+
+    Args:
+        base64_images: List[str] - List of base64 encoded images
+
+    Returns:
+        List[Optional[str]] - List of public URLs or None for failed uploads
+    """
+    tasks = []
+    for idx, base64_image in enumerate(base64_images):
+        timestamp = (
+            datetime.now(pytz.timezone("Asia/Ho_Chi_Minh"))
+            .replace(tzinfo=None)
+            .strftime("%Y-%m-%d_%H-%M-%S")
+        )
+        file_name = f"image_{timestamp}_{idx}"
+        tasks.append(upload_base64_image_to_storage(base64_image, file_name))
+
+    return await asyncio.gather(*tasks, return_exceptions=True)

src/modules/config_extractor.py

@@ -0,0 +1,66 @@
+import torchvision
+
+
+# Config for the models that are supported by the extractor
+MODEL_CONFIG = {
+    "resnet18": {
+        "weights": torchvision.models.ResNet18_Weights.DEFAULT,
+        "model": torchvision.models.resnet18,
+        "feat_layer": "flatten",
+        "feat_dims": 512,
+    },
+    "resnet34": {
+        "weights": torchvision.models.ResNet34_Weights.DEFAULT,
+        "model": torchvision.models.resnet34,
+        "feat_layer": "flatten",
+        "feat_dims": 512,
+    },
+    "resnet50": {
+        "weights": torchvision.models.ResNet50_Weights.DEFAULT,
+        "model": torchvision.models.resnet50,
+        "feat_layer": "flatten",
+        "feat_dims": 2048,
+    },
+    "resnet101": {
+        "weights": torchvision.models.ResNet101_Weights.DEFAULT,
+        "model": torchvision.models.resnet101,
+        "feat_layer": "flatten",
+        "feat_dims": 2048,
+    },
+    "resnet152": {
+        "weights": torchvision.models.ResNet152_Weights.DEFAULT,
+        "model": torchvision.models.resnet152,
+        "feat_layer": "flatten",
+        "feat_dims": 2048,
+    },
+    "vit_b_16": {
+        "weights": torchvision.models.ViT_B_16_Weights.DEFAULT,
+        "model": torchvision.models.vit_b_16,
+        "feat_layer": "getitem_5",
+        "feat_dims": 768,
+    },
+    "vit_b_32": {
+        "weights": torchvision.models.ViT_B_32_Weights.DEFAULT,
+        "model": torchvision.models.vit_b_32,
+        "feat_layer": "getitem_5",
+        "feat_dims": 768,
+    },
+    "vit_l_16": {
+        "weights": torchvision.models.ViT_L_16_Weights.DEFAULT,
+        "model": torchvision.models.vit_l_16,
+        "feat_layer": "getitem_5",
+        "feat_dims": 1024,
+    },
+    "vit_l_32": {
+        "weights": torchvision.models.ViT_L_32_Weights.DEFAULT,
+        "model": torchvision.models.vit_l_32,
+        "feat_layer": "getitem_5",
+        "feat_dims": 1024,
+    },
+    "vit_h_14": {
+        "weights": torchvision.models.ViT_H_14_Weights.DEFAULT,
+        "model": torchvision.models.vit_h_14,
+        "feat_layer": "getitem_5",
+        "feat_dims": 1280,
+    },
+}

src/modules/feature_extractor.py

@@ -0,0 +1,139 @@
+import torchvision.models.feature_extraction
+import torchvision
+import os
+import torch
+import onnx
+import onnxruntime
+
+from src.modules.config_extractor import MODEL_CONFIG
+
+os.environ["KMP_DUPLICATE_LIB_OK"] = "True"
+
+
+class FeatureExtractor:
+    """Class for extracting features from images using a pre-trained model"""
+
+    def __init__(self, base_model, onnx_path=None):
+        # set the base model
+        self.base_model = base_model
+        # get the number of features
+        self.feat_dims = MODEL_CONFIG[base_model]["feat_dims"]
+        # get the feature layer name
+        self.feat_layer = MODEL_CONFIG[base_model]["feat_layer"]
+
+        # Set default ONNX path if not provided
+        if onnx_path is None:
+            onnx_path = f"model/{base_model}_feature_extractor.onnx"
+
+        self.onnx_path = onnx_path
+        self.onnx_session = None
+
+        # Initialize transforms (needed for both ONNX and PyTorch)
+        _, self.transforms = self.init_model(base_model)
+
+        # Check if ONNX model exists
+        if os.path.exists(onnx_path):
+            print(f"Loading existing ONNX model from {onnx_path}")
+            self.onnx_session = onnxruntime.InferenceSession(onnx_path)
+        else:
+            print(
+                f"ONNX model not found at {onnx_path}. Initializing PyTorch model and converting to ONNX..."
+            )
+            # Initialize PyTorch model
+            self.model, _ = self.init_model(base_model)
+            self.model.eval()
+            self.device = torch.device("cpu")
+            self.model.to(self.device)
+
+            # Create directory if it doesn't exist
+            os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
+
+            # Convert to ONNX
+            self.convert_to_onnx(onnx_path)
+
+            # Load the newly created ONNX model
+            self.onnx_session = onnxruntime.InferenceSession(onnx_path)
+            print(f"Successfully created and loaded ONNX model from {onnx_path}")
+
+    def init_model(self, base_model):
+        """Initialize the model for feature extraction
+
+        Args:
+            base_model: str, the name of the base model
+
+        Returns:
+            model: torch.nn.Module, the feature extraction model
+            transforms: torchvision.transforms.Compose, the image transformations
+        """
+        if base_model not in MODEL_CONFIG:
+            raise ValueError(f"Invalid base model: {base_model}")
+
+        # get the model and weights
+        weights = MODEL_CONFIG[base_model]["weights"]
+        model = torchvision.models.feature_extraction.create_feature_extractor(
+            MODEL_CONFIG[base_model]["model"](weights=weights),
+            [MODEL_CONFIG[base_model]["feat_layer"]],
+        )
+        # get the image transformations
+        transforms = weights.transforms()
+        return model, transforms
+
+    def extract_features(self, img):
+        """Extract features from an image
+
+        Args:
+            img: PIL.Image, the input image
+
+        Returns:
+            output: torch.Tensor, the extracted features
+        """
+        # apply transformations
+        x = self.transforms(img)
+        # add batch dimension
+        x = x.unsqueeze(0)
+
+        # Convert to numpy for ONNX Runtime
+        x_numpy = x.numpy()
+        # Run inference with ONNX Runtime
+        print("Running inference with ONNX Runtime")
+        output = self.onnx_session.run(
+            None,
+            {'input': x_numpy}
+        )[0]
+        # Convert back to torch tensor
+        output = torch.from_numpy(output)
+
+        return output
+
+    def convert_to_onnx(self, save_path):
+        """Convert the model to ONNX format and save it
+
+        Args:
+            save_path: str, the path to save the ONNX model
+
+        Returns:
+            None
+        """
+        # Create a dummy input tensor
+        dummy_input = torch.randn(1, 3, 224, 224, device=self.device)
+
+        # Export the model
+        torch.onnx.export(
+            self.model,
+            dummy_input,
+            save_path,
+            export_params=True,
+            opset_version=14,
+            do_constant_folding=True,
+            input_names=['input'],
+            output_names=['output'],
+            dynamic_axes={
+                'input': {0: 'batch_size'},
+                'output': {0: 'batch_size'}
+            }
+        )
+
+        # Verify the exported model
+        onnx_model = onnx.load(save_path)
+        onnx.checker.check_model(onnx_model)
+        print(f"ONNX model saved to {save_path}")

src/search_query.py

@@ -0,0 +1,154 @@
+# Description:
+#   This script is used to query the index for similar images to a set of random images.
+#   The script uses the FeatureExtractor class to extract the features from the images and the Faiss index to search for similar images.
+#
+# Usage:
+#
+#   To use this script, you can run the following commands: (You MUST define a feat_extractor since indexings are different for each model)
+#       python3 search_query.py --feat_extractor resnet50
+#       python3 search_query.py --feat_extractor resnet101
+#       python3 search_query.py --feat_extractor resnet50 --n 5
+#       python3 search_query.py --feat_extractor resnet50 --k 20
+#       python3 search_query.py --feat_extractor resnet50 --n 10 --k 12
+#
+import matplotlib.pyplot as plt
+import numpy as np
+import argparse
+import torch
+import faiss
+import PIL
+import os
+
+from src.modules.feature_extractor import FeatureExtractor
+from src.config.settings import FEATURE_EXTRACTOR_MODELS
+from src.config.settings import DATA_DIR, IMAGES_DIR, RESULTS_DIR
+
+
+def select_random_images(n, image_list):
+    """Select n random images from the image list.
+
+    Args:
+        n (int): The number of images to select.
+        image_list (list[str]): The list of image file names.
+
+    Returns:
+        list[PIL.Image]: The list of selected images.
+    """
+    selected_indices = np.random.randint(len(image_list), size=n)
+    img_filenames = [image_list[i] for i in selected_indices]
+    images = [
+        PIL.Image.open(os.path.join(IMAGES_DIR, img_filename))
+        for img_filename in img_filenames
+    ]
+    return images
+
+
+def plot_query_results(query_img, similar_imgs, distances, out_filepath):
+    """Plot the query image and the similar images side by side. Save the plot to the specified file path.
+
+    Args:
+        query_img (PIL.Image): The query image.
+        similar_imgs (list[PIL.Image]): The list of similar images.
+        distances (list[float]): The list of distances of the similar images.
+        out_filepath (str): The file path to save the plot.
+
+    Returns:
+        None
+    """
+    # initialize the figure
+    fig, axes = plt.subplots(3, args.k // 2, figsize=(20, 10))
+    # plot the query image
+    axes[0, 0].imshow(query_img)
+    axes[0, 0].set_title("Query Image")
+    axes[0, 0].axis("off")
+    # do not draw the remaining pots in the first row
+    for i in range(1, args.k // 2):
+        axes[0, i].axis("off")
+    # plot the similar images
+    for i, (img, dist) in enumerate(zip(similar_imgs, distances)):
+        axes[i // (args.k // 2) + 1, i % (args.k // 2)].imshow(img)
+        axes[i // (args.k // 2) + 1, i % (args.k // 2)].set_title(f"{dist:.4f}")
+        axes[i // (args.k // 2) + 1, i % (args.k // 2)].axis("off")
+    # remove the remaining axes
+    plt.tight_layout()
+    # save the plot
+    plt.savefig(out_filepath, bbox_inches="tight", dpi=200)
+
+
+def main(args=None):
+
+    # set the random seed for reproducibility
+    np.random.seed(args.seed)
+
+    # load the vector database index
+    index_filepath = os.path.join(DATA_DIR, f"db_{args.feat_extractor}.index")
+    index = faiss.read_index(index_filepath)
+
+    # initialize the feature extractor with the base model specified in the arguments
+    feature_extractor = FeatureExtractor(base_model=args.feat_extractor)
+
+    # get the list of images in sorted order since the index is built in the same order
+    image_list = sorted(os.listdir(IMAGES_DIR))
+    # select n random images
+    query_images = select_random_images(args.n, image_list)
+
+    with torch.no_grad():
+        # iterate over the selected/query images
+        for query_idx, img in enumerate(query_images, start=1):
+            # output now has the features corresponding to input x
+            output = feature_extractor.extract_ưfeatures(img)
+            # keep only batch dimension
+            output = output.view(output.size(0), -1)
+            # normalize
+            output = output / output.norm(p=2, dim=1, keepdim=True)
+            # search for similar images
+            D, I = index.search(output.cpu().numpy(), args.k)
+
+            # get the similar images
+            similar_images = [
+                PIL.Image.open(os.path.join(IMAGES_DIR, image_list[index]))
+                for index in I[0]
+            ]
+            # plot the query results and save the plot
+            query_results_folderpath = f"{RESULTS_DIR}/results_{args.feat_extractor}"
+            os.makedirs(query_results_folderpath, exist_ok=True)
+            query_results_filepath = (
+                f"{query_results_folderpath}/query_{query_idx:03}.jpg"
+            )
+            plot_query_results(
+                img, similar_images, D[0], out_filepath=query_results_filepath
+            )
+
+
+if __name__ == "__main__":
+    # parse arguments
+    args = argparse.ArgumentParser()
+    args.add_argument(
+        "--feat_extractor",
+        type=str,
+        choices=FEATURE_EXTRACTOR_MODELS,
+        required=True,
+    )
+    args.add_argument(
+        "--n",
+        type=int,
+        default=10,
+        help="Number of random images to select",
+    )
+    args.add_argument(
+        "--k",
+        type=int,
+        default=12,
+        help="Number of similar images to retrieve",
+    )
+    args.add_argument(
+        "--seed",
+        type=int,
+        default=777,
+        help="Random seed for reproducibility",
+    )
+
+    args = args.parse_args()
+
+    # run the main function
+    main(args)

src/utils/helper.py

@@ -0,0 +1,19 @@
+import base64
+from PIL import Image
+from fastapi import HTTPException
+from io import BytesIO
+
+
+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")