app.py

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)