from fastapi import FastAPI, File, UploadFile, Response, HTTPException from fastapi.responses import JSONResponse, FileResponse, StreamingResponse from fastapi.middleware.cors import CORSMiddleware from PIL import Image from datetime import datetime import io from io import BytesIO import requests import sqlite3 from pydantic import BaseModel, EmailStr from typing import List, Optional from pathlib import Path from model import YOLOModel import shutil from openpyxl import Workbook from openpyxl.drawing.image import Image as ExcelImage from openpyxl.styles import Alignment import os yolo = YOLOModel() UPLOAD_FOLDER = Path("./uploads") UPLOAD_FOLDER.mkdir(exist_ok=True) app = FastAPI() cropped_images_dir = "cropped_images" # Initialize SQLite database def init_db(): conn = sqlite3.connect('users.db') c = conn.cursor() c.execute(''' CREATE TABLE IF NOT EXISTS users ( id INTEGER PRIMARY KEY AUTOINCREMENT, firstName TEXT NOT NULL, lastName TEXT NOT NULL, country TEXT, number TEXT, -- Phone number stored as TEXT to allow various formats email TEXT UNIQUE NOT NULL, -- Email should be unique and non-null password TEXT NOT NULL -- Password will be stored as a string (hashed ideally) ) ''') conn.commit() conn.close() init_db() class UserSignup(BaseModel): firstName: str lastName: str country: str number: str email: EmailStr password: str class UserLogin(BaseModel): email: str password: str app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) @app.post("/signup") async def signup(user_data: UserSignup): try: conn = sqlite3.connect('users.db') c = conn.cursor() # Check if user already exists c.execute("SELECT * FROM users WHERE email = ?", (user_data.email,)) if c.fetchone(): raise HTTPException(status_code=400, detail="Email already registered") # Insert new user c.execute(""" INSERT INTO users (firstName, lastName, country, number, email, password) VALUES (?, ?, ?, ?, ?, ?) """, (user_data.firstName, user_data.lastName, user_data.country, user_data.number, user_data.email, user_data.password)) conn.commit() conn.close() return {"message": "User registered successfully", "email": user_data.email} except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.post("/login") async def login(user_data: UserLogin): try: conn = sqlite3.connect('users.db') c = conn.cursor() # Find user c.execute("SELECT * FROM users WHERE email = ? AND password = ?", (user_data.email, user_data.password)) user = c.fetchone() conn.close() if not user: raise HTTPException(status_code=401, detail="Invalid credentials") return { "message": "Login successful", "user": { "firstName": user[1], "lastName": user[2], "email": user[3] } } except Exception as e: raise HTTPException(status_code=500, detail=str(e)) @app.post("/upload") async def upload_image(image: UploadFile = File(...)): # print(f'\n\t\tUPLOADED!!!!') try: file_path = UPLOAD_FOLDER / image.filename with file_path.open("wb") as buffer: shutil.copyfileobj(image.file, buffer) # print(f'Starting to pass into model, {file_path}') # Perform YOLO inference predictions = yolo.predict(str(file_path)) print(f'\n\n\n{predictions}\n\n\ \n\t\t\t\tare predictions') # Clean up uploaded file file_path.unlink() # Remove file after processing return JSONResponse(content={"items": predictions}) except Exception as e: return JSONResponse(content={"error": str(e)}, status_code=500) @app.get("/download_cropped_image/{image_idx}") def download_cropped_image(image_idx: int): cropped_image_path = cropped_images_dir / f"crop_{image_idx}.jpg" if cropped_image_path.exists(): return FileResponse(cropped_image_path, media_type="image/jpeg") return JSONResponse(content={"error": "Cropped image not found"}, status_code=404) def cleanup_images(directory: str): """Remove all images in the directory.""" for file in Path(directory).glob("*"): file.unlink() ''' @app.post("/generate-excel/") async def generate_excel(predictions: list): # Create an Excel workbook workbook = Workbook() sheet = workbook.active sheet.title = "Predictions" # Add headers headers = ["Category", "Confidence", "Predicted Brand", "Price", "Details", "Detected Text", "Image"] sheet.append(headers) for idx, prediction in enumerate(predictions): # Extract details from the prediction category = prediction["category"] confidence = prediction["confidence"] predicted_brand = prediction["predicted_brand"] price = prediction["price"] details = prediction["details"] detected_text = prediction["detected_text"] cropped_image_path = prediction["image_path"] # Append data row sheet.append([category, confidence, predicted_brand, price, details, detected_text]) # Add the image to the Excel file (if it exists) if os.path.exists(cropped_image_path): img = ExcelImage(cropped_image_path) img.width, img.height = 50, 50 # Resize image to fit into the cell sheet.add_image(img, f"G{idx + 2}") # Place in the "Image" column excel_file_path = "predictions_with_images.xlsx" workbook.save(excel_file_path) # Cleanup after saving cleanup_images(cropped_images_dir) # Serve the Excel file as a response return FileResponse( excel_file_path, media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", filename="predictions_with_images.xlsx" ) ''' # Define the Prediction model class Prediction(BaseModel): category: Optional[str] confidence: Optional[float] predicted_brand: Optional[str] price: Optional[str] details: Optional[str] detected_text: Optional[str] image_url: Optional[str] image_path: Optional[str] @app.post("/generate-excel/") async def generate_excel(predictions: List[Prediction]): print('Generate excel called') # Create an Excel workbook workbook = Workbook() sheet = workbook.active sheet.title = "Predictions" # Add headers headers = ["Category", "Confidence", "Predicted Brand", "Price", "Image URL", "Details", "Detected Text"] sheet.append(headers) # Add prediction rows (skipping for brevity) # Set header style and alignment for cell in sheet[1]: cell.alignment = Alignment(horizontal="center", vertical="center", wrap_text=True) sheet.row_dimensions[1].height = 30 # Adjust header row height # Set column widths based on data type column_widths = { "A": 20, # Category "B": 15, # Confidence "C": 40, # Predicted Brand "D": 15, # Price "E": 50, # Image URL "F": 30, # Details "G": 30 # Detected Text } for col, width in column_widths.items(): sheet.column_dimensions[col].width = width # Add prediction rows for idx, prediction in enumerate(predictions): row_index = idx + 2 # Start from the second row # Add data to the row sheet.append([ prediction.category, prediction.confidence, prediction.predicted_brand, prediction.price, prediction.image_url, prediction.details, prediction.detected_text, ]) # Adjust row height for multiline text sheet.row_dimensions[row_index].height = 180 # Default height for rows # Wrap text in all cells of the row for col_idx in range(1, 8): # Columns A to G cell = sheet.cell(row=row_index, column=col_idx) cell.alignment = Alignment(wrap_text=True, vertical="top") if prediction.image_url: try: response = requests.get(prediction.image_url) img = ExcelImage(BytesIO(response.content)) img.width, img.height = 160, 160 # Resize image to fit into the cell img_cell = f"G{row_index}" # Image column sheet.add_image(img, img_cell) except requests.exceptions.RequestException as e: print(f"Error downloading image: {e}") # # Add image if the path exists # if os.path.exists(prediction.image_path): # img = ExcelImage(prediction.image_path) # img.width, img.height = 160, 160 # Resize image to fit into the cell # img_cell = f"G{row_index}" # Image column # sheet.add_image(img, img_cell) # Create a unique filename based on the current timestamp or index timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") excel_file_path = f"/predictions_with_images_{timestamp}.xlsx" print(excel_file_path) # Save the Excel file to the specified path workbook.save(excel_file_path) # Check if the directory exists, if not, create it (to store multiple files) if not os.path.exists("/predictions"): os.makedirs("/predictions") # Move the file to a new directory os.rename(excel_file_path, f"/predictions/{os.path.basename(excel_file_path)}") hf_path = "https://huggingface.co/spaces/root-sajjan/whatisit/resolve/main" excel_file_path = hf_path + f"/predictions/{os.path.basename(excel_file_path)}"+"?download=True" return JSONResponse(content={"download_link": excel_file_path}) # else: # return JSONResponse(status_code=500, content={"error": "File upload failed"}) ''' @app.post("/generate-excel/") async def generate_excel(predictions: List[Prediction]): print('Generate excel called') # Create an Excel workbook workbook = Workbook() sheet = workbook.active sheet.title = "Predictions" # Add headers headers = ["Category", "Confidence", "Predicted Brand", "Price", "Image URL", "Details", "Detected Text"] sheet.append(headers) # Set header style and alignment for cell in sheet[1]: cell.alignment = Alignment(horizontal="center", vertical="center", wrap_text=True) sheet.row_dimensions[1].height = 30 # Adjust header row height # Set column widths based on data type column_widths = { "A": 20, # Category "B": 15, # Confidence "C": 40, # Predicted Brand "D": 15, # Price "E": 50, # Image URL "F": 30, # Details "G": 30 # Detected Text } for col, width in column_widths.items(): sheet.column_dimensions[col].width = width # Add prediction rows for idx, prediction in enumerate(predictions): row_index = idx + 2 # Start from the second row # Add data to the row sheet.append([ prediction.category, prediction.confidence, prediction.predicted_brand, prediction.price, prediction.image_url, prediction.details, prediction.detected_text, ]) # Adjust row height for multiline text sheet.row_dimensions[row_index].height = 180 # Default height for rows # Wrap text in all cells of the row for col_idx in range(1, 8): # Columns A to G cell = sheet.cell(row=row_index, column=col_idx) cell.alignment = Alignment(wrap_text=True, vertical="top") # If image URL is provided, download it if prediction.image_url: try: response = requests.get(prediction.image_url) img = ExcelImage(BytesIO(response.content)) img.width, img.height = 160, 160 # Resize image to fit into the cell img_cell = f"G{row_index}" # Image column sheet.add_image(img, img_cell) except requests.exceptions.RequestException as e: print(f"Error downloading image: {e}") # Optionally add a placeholder image or text # Save the Excel file excel_file_path = "/tmp/predictions_with_images.xlsx" workbook.save(excel_file_path) # Serve the Excel file as a response return FileResponse( excel_file_path, media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", filename="predictions_with_images.xlsx" ) @app.post("/generate-excel2/") async def generate_excel(predictions: List[Prediction]): print('Generate excel called') # Create an Excel workbook workbook = Workbook() sheet = workbook.active sheet.title = "Predictions" # Add headers headers = ["Category", "Confidence", "Predicted Brand", "Price", "Image URL", "Details", "Detected Text", ] sheet.append(headers) # Set header style and alignment for cell in sheet[1]: cell.alignment = Alignment(horizontal="center", vertical="center", wrap_text=True) sheet.row_dimensions[1].height = 30 # Adjust header row height # Set column widths based on data type column_widths = { "A": 20, # Category "B": 15, # Confidence "C": 40, # Predicted Brand "D": 15, # Price "E": 50, # Image URL "F": 30, # Details "G": 30 # Detected Text } for col, width in column_widths.items(): sheet.column_dimensions[col].width = width # Add prediction rows for idx, prediction in enumerate(predictions): row_index = idx + 2 # Start from the second row # Add data to the row sheet.append([ prediction.category, prediction.confidence, prediction.predicted_brand, prediction.price, prediction.image_url, prediction.details, prediction.detected_text, ]) # Adjust row height for multiline text sheet.row_dimensions[row_index].height = 180 # Default height for rows # Wrap text in all cells of the row for col_idx in range(1, 8): # Columns A to G cell = sheet.cell(row=row_index, column=col_idx) cell.alignment = Alignment(wrap_text=True, vertical="top") # Add image if the path exists if os.path.exists(prediction.image_path): img = ExcelImage(prediction.image_path) img.width, img.height = 160, 160 # Resize image to fit into the cell img_cell = f"G{row_index}" # Image column sheet.add_image(img, img_cell) # Save the Excel file excel_file_path = "predictions_with_images.xlsx" workbook.save(excel_file_path) # Serve the Excel file as a response return FileResponse( excel_file_path, media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", filename="predictions_with_images.xlsx" ) ''' ''' @app.post("/generate-excel/") async def generate_excel(predictions: list): print('Generate excel called') # Create an Excel workbook workbook = Workbook() sheet = workbook.active sheet.title = "Predictions" # Add headers headers = ["Category", "Confidence", "Predicted Brand", "Price", "Details", "Detected Text", "Image URL"] sheet.append(headers) # Format the header row for cell in sheet[1]: cell.alignment = Alignment(horizontal="center", vertical="center") for idx, prediction in enumerate(predictions): # Extract details from the prediction category = prediction["category"] confidence = prediction["confidence"] predicted_brand = prediction["predicted_brand"] price = prediction["price"] details = prediction["details"] detected_text = prediction["detected_text"] image_url = prediction["image_url"] # URL to the image cropped_image_path = prediction["image_path"] # Path to local image file for Excel embedding # Append data row sheet.append([category, confidence, predicted_brand, price, details, detected_text, image_url]) # If the image path exists, add the image to the Excel file if os.path.exists(cropped_image_path): img = ExcelImage(cropped_image_path) img.width, img.height = 50, 50 # Resize image to fit into the cell sheet.add_image(img, f"G{idx + 2}") # Place in the "Image" column excel_file_path = "predictions_with_images.xlsx" workbook.save(excel_file_path) # Serve the Excel file as a response return FileResponse( excel_file_path, media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", filename="predictions_with_images.xlsx" ) ''' # code to accept the localhost to get images from if __name__ == "__main__": import uvicorn uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)