updated

main.py

@@ -1,333 +1,337 @@
-from fastapi import FastAPI, File, UploadFile, Response, HTTPException
-from fastapi.responses import JSONResponse, FileResponse
-from fastapi.middleware.cors import CORSMiddleware
-from PIL import Image
-import io
-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.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)
-
-    # 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
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["http://192.168.56.1:3000", "http://192.168.56.1:3001", "http://localhost:3000"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+from fastapi import FastAPI, File, UploadFile, Response, HTTPException
+from fastapi.responses import JSONResponse, FileResponse
+from fastapi.middleware.cors import CORSMiddleware
+from PIL import Image
+import io
+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.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)
+
+    # 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
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)